feat: add non-ZeroGPU endpoints, automatic failover, and endpoint cooldown

qflash_test.go

@@ -232,7 +232,7 @@ func TestQwenServiceChatMock(t *testing.T) {
   }))
   defer mockServer.Close()
 
-  svc := NewQwenService(mockServer.URL, "Qwen/Qwen3.8-Flash-Next", "chat_response", "", "", "", "", true)
+  svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-Flash-Next", "chat_response", "", "", "", "", true, true)
 
   // 1. Test Non-streaming completion
   rec := httptest.NewRecorder()
@@ -362,7 +362,7 @@ func TestQwenServiceChatRespondMock(t *testing.T) {
   }))
   defer mockServer.Close()
 
-  svc := NewQwenService(mockServer.URL, "Qwen/Qwen3.8-27B-Uncensored", "respond", "", "", "", "", true)
+  svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "respond", "", "", "", "", true, true)
 
   // Non-streaming test
   rec := httptest.NewRecorder()
@@ -445,3 +445,185 @@ func TestQwenServiceChatRespondMock(t *testing.T) {
   }
 }
 
+func TestQwenServiceDirectOpenAIMock(t *testing.T) {
+  mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    if r.URL.Path == "/v1/chat/completions" && r.Method == http.MethodPost {
+      var req ChatCompletionRequest
+      if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+        http.Error(w, "bad request", http.StatusBadRequest)
+        return
+      }
+
+      if !req.Stream {
+        w.Header().Set("Content-Type", "application/json")
+        json.NewEncoder(w).Encode(ChatCompletionResponse{
+          ID:      "chatcmpl-openai-mock",
+          Object:  "chat.completion",
+          Created: 1234567890,
+          Model:   req.Model,
+          Choices: []ChatCompletionResponseChoice{
+            {
+              Index: 0,
+              Message: ChatMessage{
+                Role:    "assistant",
+                Content: "Hello from native OpenAI upstream!",
+              },
+              FinishReason: "stop",
+            },
+          },
+        })
+        return
+      }
+
+      // Stream
+      w.Header().Set("Content-Type", "text/event-stream")
+      flusher, _ := w.(http.Flusher)
+      chunk := `data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"Hello from "}}]}
+
+data: {"id":"chatcmpl-openai-mock","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"content":"native stream!"}}]}
+
+data: [DONE]
+
+`
+      w.Write([]byte(chunk))
+      if flusher != nil {
+        flusher.Flush()
+      }
+      return
+    }
+
+    http.NotFound(w, r)
+  }))
+  defer mockServer.Close()
+
+  svc := NewQwenService([]string{mockServer.URL}, "Qwen/Qwen3.8-27B-Uncensored", "openai", "", "", "", "", true, true)
+
+  // Non-streaming
+  rec := httptest.NewRecorder()
+  req := ChatCompletionRequest{
+    Model: "qwen-openai",
+    Messages: []ChatMessage{
+      {Role: "user", Content: "Hello"},
+    },
+    Stream: false,
+  }
+  if err := svc.Chat(rec, nil, req); err != nil {
+    t.Fatalf("unexpected error in direct openai non-streaming: %v", err)
+  }
+  if rec.Code != http.StatusOK {
+    t.Fatalf("expected HTTP 200, got %d", rec.Code)
+  }
+  var resp ChatCompletionResponse
+  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
+    t.Fatalf("failed to decode response: %v", err)
+  }
+  if resp.Choices[0].Message.Content != "Hello from native OpenAI upstream!" {
+    t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
+  }
+
+  // Streaming
+  recStream := httptest.NewRecorder()
+  reqStream := ChatCompletionRequest{
+    Model: "qwen-openai",
+    Messages: []ChatMessage{
+      {Role: "user", Content: "Hello"},
+    },
+    Stream: true,
+  }
+  if err := svc.Chat(recStream, nil, reqStream); err != nil {
+    t.Fatalf("unexpected error in direct openai streaming: %v", err)
+  }
+  if !strings.Contains(recStream.Body.String(), "native stream!") {
+    t.Fatalf("expected streamed content, got %q", recStream.Body.String())
+  }
+}
+
+func TestQwenServiceAutoFailover(t *testing.T) {
+  // Server 1 simulates ZeroGPU quota exceeded (returns HTTP 429)
+  server1Hits := 0
+  server1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    server1Hits++
+    w.WriteHeader(http.StatusTooManyRequests)
+    w.Write([]byte(`{"error": "ZeroGPU runs limit reached: quota exceeded"}`))
+  }))
+  defer server1.Close()
+
+  // Server 2 is healthy and serves completions
+  server2Hits := 0
+  server2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    server2Hits++
+    if r.URL.Path == "/v1/chat/completions" {
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(ChatCompletionResponse{
+        ID:      "chatcmpl-failover-mock",
+        Object:  "chat.completion",
+        Created: 1234567890,
+        Choices: []ChatCompletionResponseChoice{
+          {
+            Index: 0,
+            Message: ChatMessage{
+              Role:    "assistant",
+              Content: "Response from fallback server!",
+            },
+            FinishReason: "stop",
+          },
+        },
+      })
+      return
+    }
+    http.NotFound(w, r)
+  }))
+  defer server2.Close()
+
+  endpoints := []string{server1.URL, server2.URL}
+  svc := NewQwenService(endpoints, "Qwen/Qwen3.8-27B", "openai", "", "", "", "", true, true)
+
+  rec := httptest.NewRecorder()
+  req := ChatCompletionRequest{
+    Model: "qwen3.8-27b",
+    Messages: []ChatMessage{
+      {Role: "user", Content: "Failover test"},
+    },
+    Stream: false,
+  }
+
+  err := svc.Chat(rec, nil, req)
+  if err != nil {
+    t.Fatalf("expected failover to succeed, got error: %v", err)
+  }
+
+  if rec.Code != http.StatusOK {
+    t.Fatalf("expected HTTP 200 after failover, got %d", rec.Code)
+  }
+
+  var resp ChatCompletionResponse
+  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
+    t.Fatalf("failed to decode failover response: %v", err)
+  }
+
+  if resp.Choices[0].Message.Content != "Response from fallback server!" {
+    t.Fatalf("unexpected content: %v", resp.Choices[0].Message.Content)
+  }
+
+  if server1Hits != 1 {
+    t.Fatalf("expected server1 to be attempted once, got %d", server1Hits)
+  }
+  if server2Hits != 1 {
+    t.Fatalf("expected server2 to be called once on failover, got %d", server2Hits)
+  }
+
+  // Second request should skip server 1 because it entered cooldown
+  rec2 := httptest.NewRecorder()
+  err2 := svc.Chat(rec2, nil, req)
+  if err2 != nil {
+    t.Fatalf("expected second request to succeed directly on server2: %v", err2)
+  }
+  if server1Hits != 1 {
+    t.Fatalf("expected server1 to be skipped during cooldown, but was hit %d times", server1Hits)
+  }
+  if server2Hits != 2 {
+    t.Fatalf("expected server2 to receive the second request, got hits: %d", server2Hits)
+  }
+}
+
+