Restrict call_v2 to explicit code snippets and omit session_hash from call payloads

This commit is contained in:
Luxferre
2026-09-07 16:14:26 +03:00
parent 36bbe2bf86
commit 2a0fd63676
2 changed files with 140 additions and 15 deletions
+131 -5
View File
@@ -2101,9 +2101,9 @@ func TestInspectSpaceBlocksChatbotStateResolution(t *testing.T) {
}
}
// TestExecuteCallCompletionSessionHashAndTrailingStateTrim verifies that call requests
// send a valid session_hash and trim unexposed trailing state inputs from the payload.
func TestExecuteCallCompletionSessionHashAndTrailingStateTrim(t *testing.T) {
// TestExecuteCallCompletionTrailingStateTrim verifies that call requests
// omit session_hash (to avoid breaking SSE streams) and trim unexposed trailing state inputs from the payload.
func TestExecuteCallCompletionTrailingStateTrim(t *testing.T) {
var receivedSessionHash string
var receivedDataLen int
@@ -2204,8 +2204,8 @@ func TestExecuteCallCompletionSessionHashAndTrailingStateTrim(t *testing.T) {
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 receivedSessionHash != "" {
t.Errorf("expected empty session_hash in call payload, got %s", receivedSessionHash)
}
if receivedDataLen != 1 {
t.Errorf("expected call data length 1 (unexposed state trimmed), got %d", receivedDataLen)
@@ -2337,6 +2337,132 @@ func TestGradio6CallV2Protocol(t *testing.T) {
}
}
// TestGradio6WithoutSnippetResolvesToCall verifies that Gradio 6 spaces without code snippets
// resolve protocol call (not call_v2), submit standard data array, and omit session_hash.
func TestGradio6WithoutSnippetResolvesToCall(t *testing.T) {
var callEndpointReached string
var receivedPayload map[string]interface{}
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{}{
"/predict": map[string]interface{}{
"parameters": []map[string]interface{}{
{
"parameter_name": "message",
"component": "Textbox",
"label": "",
},
},
// No code_snippets present
},
},
})
return
}
if r.URL.Path == "/config" {
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"version": "6.5.1",
"components": []map[string]interface{}{
{"id": 11, "type": "textbox", "props": map[string]interface{}{"label": ""}},
{"id": 15, "type": "state"},
},
"dependencies": []map[string]interface{}{
{
"id": 6,
"api_name": "predict",
"inputs": []int{11, 15},
"outputs": []int{11, 15},
"types": map[string]interface{}{"generator": true},
},
},
})
return
}
if r.URL.Path == "/gradio_api/call/predict" {
callEndpointReached = r.URL.Path
if err := json.NewDecoder(r.Body).Decode(&receivedPayload); err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
w.Header().Set("Content-Type", "application/json")
json.NewEncoder(w).Encode(map[string]interface{}{
"event_id": "evt-call-651",
})
return
}
if r.URL.Path == "/gradio_api/call/predict/evt-call-651" {
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: complete\ndata: [\"Hello from standard Gradio 6 call!\"]\n\n")
if flusher != nil {
flusher.Flush()
}
return
}
http.NotFound(w, r)
}))
defer ts.Close()
gw := NewGradioGateway(ts.URL, "", 10*time.Second)
// 1. Verify inspection resolved protocol "call", not "call_v2"
disc, err := InspectSpace(gw.client, ts.URL, DefaultUserAgent)
if err != nil {
t.Fatalf("InspectSpace failed: %v", err)
}
if disc.Protocol != "call" {
t.Errorf("expected Protocol call, got %s", disc.Protocol)
}
if disc.Endpoint != "/predict" {
t.Errorf("expected Endpoint /predict, got %s", disc.Endpoint)
}
if disc.TotalInputs != 1 {
t.Errorf("expected TotalInputs 1 (trimmed trailing state), got %d", disc.TotalInputs)
}
if disc.RawTotalInputs != 2 {
t.Errorf("expected RawTotalInputs 2, got %d", disc.RawTotalInputs)
}
// 2. Execute chat completion
req := ChatCompletionRequest{
Model: "default",
Messages: []ChatMessage{
{Role: "user", Content: "Hello test"},
},
}
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 callEndpointReached != "/gradio_api/call/predict" {
t.Errorf("expected call to /gradio_api/call/predict, got %s", callEndpointReached)
}
if _, hasHash := receivedPayload["session_hash"]; hasHash {
t.Errorf("expected no session_hash in call payload")
}
dataSlice, ok := receivedPayload["data"].([]interface{})
if !ok || len(dataSlice) != 1 || dataSlice[0] != "Hello test" {
t.Errorf("expected data payload [\"Hello test\"], got %v", receivedPayload["data"])
}
}
// TestGradioQueueProtocol verifies direct queue join and queue data streaming.
func TestGradioQueueProtocol(t *testing.T) {
var joinReached bool