diff --git a/gr2gw.go b/gr2gw.go index 6335560..1b6a1d7 100644 --- a/gr2gw.go +++ b/gr2gw.go @@ -2350,6 +2350,54 @@ func InspectSpace(client *http.Client, rawURL, userAgent string) (*SpaceDiscover } } } + + // If the endpoint has canonical parameters exposed via /gradio_api/info, + // and trailing inputs in bestMatchingDep.Inputs are unexposed server-side state components, + // do not pad them with null so Gradio preserves its internal server state. + if bestEndpointInfo != nil && len(bestEndpointInfo.Parameters) > 0 && len(bestMatchingDep.Inputs) > len(bestEndpointInfo.Parameters) { + allTrailingAreState := true + for i := len(bestEndpointInfo.Parameters); i < len(bestMatchingDep.Inputs); i++ { + cID := bestMatchingDep.Inputs[i] + if comp, exists := compMap[cID]; exists { + cType := strings.ToLower(comp.Type) + if cType != "state" && cType != "browserstate" { + allTrailingAreState = false + break + } + } + } + if allTrailingAreState { + discovery.TotalInputs = len(bestEndpointInfo.Parameters) + discovery.DefaultInputs = discovery.DefaultInputs[:discovery.TotalInputs] + if len(discovery.ParamMappings) > discovery.TotalInputs { + discovery.ParamMappings = discovery.ParamMappings[:discovery.TotalInputs] + } + if discovery.MessageIndex >= discovery.TotalInputs { + discovery.MessageIndex = -1 + } + if discovery.HistoryIndex >= discovery.TotalInputs { + discovery.HistoryIndex = -1 + } + if discovery.SystemIndex >= discovery.TotalInputs { + discovery.SystemIndex = -1 + } + if discovery.FunctionsJSONIndex >= discovery.TotalInputs { + discovery.FunctionsJSONIndex = -1 + } + if discovery.TempIndex >= discovery.TotalInputs { + discovery.TempIndex = -1 + } + if discovery.MaxTokensIndex >= discovery.TotalInputs { + discovery.MaxTokensIndex = -1 + } + if discovery.TopPIndex >= discovery.TotalInputs { + discovery.TopPIndex = -1 + } + if discovery.StreamIndex >= discovery.TotalInputs { + discovery.StreamIndex = -1 + } + } + } } // Refine history format or discover tools from bestEndpointInfo @@ -3415,7 +3463,10 @@ func (g *GradioGateway) ExecuteChatCompletion(w http.ResponseWriter, r *http.Req return g.executePredictCompletion(w, r, disc, gradioData, req, completionID, createdTime, modelName, effUA) } - payloadMap := map[string]interface{}{"data": gradioData} + payloadMap := map[string]interface{}{ + "data": gradioData, + "session_hash": GenerateUUID(), + } jsonPayload, err := json.Marshal(payloadMap) if err != nil { return fmt.Errorf("failed to encode request: %w", err) diff --git a/gr2gw_test.go b/gr2gw_test.go index 08588f1..3b00804 100644 --- a/gr2gw_test.go +++ b/gr2gw_test.go @@ -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) } } +