Support Gradio queue protocol with diff streaming and call fallback

gr2gw_test.go

@@ -2321,6 +2321,196 @@ func TestGradio6CallV2Protocol(t *testing.T) {
   }
 }
 
+// TestGradioQueueProtocol verifies direct queue join and queue data streaming.
+func TestGradioQueueProtocol(t *testing.T) {
+  var joinReached bool
+  var receivedFnIndex int
+
+  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    if r.URL.Path == "/gradio_api/queue/join" {
+      joinReached = true
+      var body struct {
+        Data        []interface{} `json:"data"`
+        FnIndex     int           `json:"fn_index"`
+        SessionHash string        `json:"session_hash"`
+      }
+      if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+        http.Error(w, err.Error(), http.StatusBadRequest)
+        return
+      }
+      receivedFnIndex = body.FnIndex
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "evt-q-1"})
+      return
+    }
+
+    if strings.HasPrefix(r.URL.Path, "/gradio_api/queue/data") {
+      w.Header().Set("Content-Type", "text/event-stream")
+      flusher, _ := w.(http.Flusher)
+      fmt.Fprintf(w, "data: {\"msg\":\"process_generating\",\"output\":{\"data\":[\"Hello\",null]},\"success\":true}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      fmt.Fprintf(w, "data: {\"msg\":\"process_generating\",\"output\":{\"data\":[[[\"append\",[],\" world!\"]]]},\"success\":true}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      fmt.Fprintf(w, "data: {\"msg\":\"process_completed\",\"output\":{\"data\":[\"Hello world!\",null]},\"success\":true}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      fmt.Fprintf(w, "data: {\"msg\":\"close_stream\"}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      return
+    }
+
+    http.NotFound(w, r)
+  }))
+  defer ts.Close()
+
+  gw := NewGradioGateway(ts.URL, "", 10*time.Second)
+  disc := &SpaceDiscovery{
+    SpaceURL:      ts.URL,
+    APIPrefix:     "/gradio_api",
+    Protocol:      "queue",
+    FnIndex:       7,
+    TotalInputs:   1,
+    RawTotalInputs: 2,
+    DefaultInputs: []interface{}{nil, nil},
+    ParamMappings: []SpaceParamMapping{
+      {InputIndex: 0, ParamType: "message"},
+      {InputIndex: 1, ParamType: "state"},
+    },
+  }
+  gw.discoveries[ts.URL] = disc
+  disc.LastDiscovered = time.Now()
+
+  req := ChatCompletionRequest{
+    Model:    "default",
+    Messages: []ChatMessage{{Role: "user", Content: "Hi"}},
+  }
+  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
+  rec := httptest.NewRecorder()
+
+  err := gw.ExecuteChatCompletion(rec, httpReq, req)
+  if err != nil {
+    t.Fatalf("ExecuteChatCompletion failed: %v", err)
+  }
+  if !joinReached {
+    t.Errorf("expected /gradio_api/queue/join to be reached")
+  }
+  if receivedFnIndex != 7 {
+    t.Errorf("expected fn_index 7, got %d", receivedFnIndex)
+  }
+
+  var resp ChatCompletionResponse
+  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
+    t.Fatalf("failed to decode response: %v", err)
+  }
+  if len(resp.Choices) == 0 || resp.Choices[0].Message.Content != "Hello world!" {
+    t.Errorf("expected content 'Hello world!', got %v", resp.Choices)
+  }
+}
+
+// TestGradioCallFallbackToQueue verifies that when /call returns a protocol or input binding error,
+// the gateway automatically falls back to /queue/join and completes successfully.
+func TestGradioCallFallbackToQueue(t *testing.T) {
+  var callReached, queueReached bool
+
+  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    if r.URL.Path == "/gradio_api/call/v2/lisa_stream" {
+      callReached = true
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "call-err-1"})
+      return
+    }
+
+    if r.URL.Path == "/gradio_api/call/lisa_stream/call-err-1" {
+      w.Header().Set("Content-Type", "text/event-stream")
+      flusher, _ := w.(http.Flusher)
+      fmt.Fprintf(w, "event: error\ndata: {\"error\": null}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      return
+    }
+
+    if r.URL.Path == "/gradio_api/queue/join" {
+      queueReached = true
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{"event_id": "queue-succ-1"})
+      return
+    }
+
+    if strings.HasPrefix(r.URL.Path, "/gradio_api/queue/data") {
+      w.Header().Set("Content-Type", "text/event-stream")
+      flusher, _ := w.(http.Flusher)
+      fmt.Fprintf(w, "data: {\"msg\":\"process_completed\",\"output\":{\"data\":[\"Recovered via queue!\",null]},\"success\":true}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      fmt.Fprintf(w, "data: {\"msg\":\"close_stream\"}\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      return
+    }
+
+    http.NotFound(w, r)
+  }))
+  defer ts.Close()
+
+  gw := NewGradioGateway(ts.URL, "", 10*time.Second)
+  disc := &SpaceDiscovery{
+    SpaceURL:      ts.URL,
+    APIPrefix:     "/gradio_api",
+    Endpoint:      "/lisa_stream",
+    CleanEndpoint: "lisa_stream",
+    Protocol:      "call_v2",
+    FnIndex:       7,
+    TotalInputs:   1,
+    RawTotalInputs: 2,
+    DefaultInputs: []interface{}{nil, nil},
+    ParamMappings: []SpaceParamMapping{
+      {InputIndex: 0, ParamName: "message", ParamType: "message"},
+    },
+  }
+  gw.discoveries[ts.URL] = disc
+  disc.LastDiscovered = time.Now()
+
+  req := ChatCompletionRequest{
+    Model:    "default",
+    Messages: []ChatMessage{{Role: "user", Content: "Test fallback"}},
+  }
+  httpReq := httptest.NewRequest("POST", "/v1/chat/completions", nil)
+  rec := httptest.NewRecorder()
+
+  err := gw.ExecuteChatCompletion(rec, httpReq, req)
+  if err != nil {
+    t.Fatalf("ExecuteChatCompletion failed: %v", err)
+  }
+  if !callReached {
+    t.Errorf("expected /call to be attempted first")
+  }
+  if !queueReached {
+    t.Errorf("expected fallback to /queue/join")
+  }
+
+  var resp ChatCompletionResponse
+  if err := json.NewDecoder(rec.Body).Decode(&resp); err != nil {
+    t.Fatalf("failed to decode response: %v", err)
+  }
+  if len(resp.Choices) == 0 || resp.Choices[0].Message.Content != "Recovered via queue!" {
+    t.Errorf("expected content 'Recovered via queue!', got %v", resp.Choices)
+  }
+  if disc.Protocol != "queue" {
+    t.Errorf("expected disc.Protocol to be switched to 'queue', got %s", disc.Protocol)
+  }
+}
+
+