Trim unexposed server-side state inputs and attach session_hash in Gradio call requests

2026-09-07 13:54:56 +03:00
parent 30f5eaceb0
commit 03825fecfd
2 changed files with 167 additions and 3 deletions
@@ -2079,10 +2079,123 @@ func TestInspectSpaceBlocksChatbotStateResolution(t *testing.T) {
  if disc.HistoryFormat != "gradio_messages" {
    t.Errorf("expected HistoryFormat gradio_messages, got %s", disc.HistoryFormat)
  }
-  if disc.TotalInputs != 2 {
-    t.Errorf("expected TotalInputs 2, got %d", disc.TotalInputs)
+  // TotalInputs should be trimmed to 1 to match the exposed API parameters and avoid padding unexposed state
+  if disc.TotalInputs != 1 {
+    t.Errorf("expected TotalInputs 1 (unexposed trailing state trimmed), got %d", disc.TotalInputs)
+  }
+}
+
+// TestExecuteCallCompletionSessionHashAndTrailingStateTrim verifies that call requests
+// send a valid session_hash and trim unexposed trailing state inputs from the payload.
+func TestExecuteCallCompletionSessionHashAndTrailingStateTrim(t *testing.T) {
+  var receivedSessionHash string
+  var receivedDataLen int
+
+  ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+    if r.URL.Path == "/gradio_api/info" {
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{
+        "named_endpoints": map[string]interface{}{
+          "/Chat_Message": map[string]interface{}{
+            "parameters": []map[string]interface{}{
+              {
+                "parameter_name": "history",
+                "component":      "Chatbot",
+                "type":           map[string]interface{}{"title": "ChatbotDataMessages"},
+                "python_type":    map[string]interface{}{"type": "dict(text: str)"},
+              },
+            },
+          },
+        },
+      })
+      return
+    }
+
+    if r.URL.Path == "/config" {
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{
+        "version": "6.20.0",
+        "components": []map[string]interface{}{
+          {"id": 5, "type": "state"},
+          {"id": 6, "type": "chatbot"},
+        },
+        "dependencies": []map[string]interface{}{
+          {
+            "id":       3,
+            "api_name": "Chat_Message",
+            "inputs":   []int{6, 5},
+            "outputs":  []int{6, 5},
+            "types":    map[string]interface{}{"generator": true},
+          },
+        },
+      })
+      return
+    }
+
+    if r.URL.Path == "/gradio_api/call/Chat_Message" {
+      var body struct {
+        Data        []interface{} `json:"data"`
+        SessionHash string        `json:"session_hash"`
+      }
+      if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
+        http.Error(w, err.Error(), http.StatusBadRequest)
+        return
+      }
+      receivedSessionHash = body.SessionHash
+      receivedDataLen = len(body.Data)
+      w.Header().Set("Content-Type", "application/json")
+      json.NewEncoder(w).Encode(map[string]interface{}{
+        "event_id": "evt-12345",
+      })
+      return
+    }
+
+    if r.URL.Path == "/gradio_api/call/Chat_Message/evt-12345" {
+      w.Header().Set("Content-Type", "text/event-stream")
+      w.Header().Set("Cache-Control", "no-cache")
+      w.WriteHeader(http.StatusOK)
+      flusher, _ := w.(http.Flusher)
+      fmt.Fprintf(w, "event: generating\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      fmt.Fprintf(w, "event: complete\ndata: [[{\"role\":\"user\",\"content\":\"hello\"},{\"role\":\"assistant\",\"content\":\"hi\"}]]\n\n")
+      if flusher != nil {
+        flusher.Flush()
+      }
+      return
+    }
+
+    http.NotFound(w, r)
+  }))
+  defer ts.Close()
+
+  gw := NewGradioGateway(ts.URL, "", 10*time.Second)
+  req := ChatCompletionRequest{
+    Model: "test-model",
+    Messages: []ChatMessage{
+      {Role: "user", Content: "hello"},
+    },
+  }
+  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 rec.Code != http.StatusOK {
+    t.Fatalf("expected status 200, got %d: %s", rec.Code, rec.Body.String())
+  }
+  if receivedSessionHash == "" {
+    t.Errorf("expected non-empty session_hash in call payload")
+  }
+  if receivedDataLen != 1 {
+    t.Errorf("expected call data length 1 (unexposed state trimmed), got %d", receivedDataLen)
  }
 }



+