From 06eeb95b4e050ede712ef16fb6053e4fa5fcfbde Mon Sep 17 00:00:00 2001
From: PP1 <74917296+pengpengyi92@users.noreply.github.com>
Date: Tue, 11 Aug 2026 04:20:54 +0800
Subject: [PATCH] fix(embedding): validate WeKnoraCloud response indexes

Reject out-of-range and duplicate response indexes before mapping them to batch inputs. Return an error when a requested input has no response item instead of silently passing an incomplete batch downstream.

Add focused coverage for malformed responses and valid out-of-order results.

Signed-off-by: PP1 <74917296+pengpengyi92@users.noreply.github.com>
---
 internal/models/embedding/weknoracloud.go     |  19 +++-
 .../models/embedding/weknoracloud_test.go     | 100 ++++++++++++++++++
 2 files changed, 117 insertions(+), 2 deletions(-)
 create mode 100644 internal/models/embedding/weknoracloud_test.go

diff --git a/internal/models/embedding/weknoracloud.go b/internal/models/embedding/weknoracloud.go
index d2285056b..f5c855442 100644
--- a/internal/models/embedding/weknoracloud.go
+++ b/internal/models/embedding/weknoracloud.go
@@ -130,9 +130,24 @@ func (e *WeKnoraCloudEmbedder) BatchEmbed(ctx context.Context, texts []string) (
 	}
 
 	result := make([][]float32, len(texts))
+	seen := make([]bool, len(texts))
 	for _, item := range embedResp.Data {
-		if item.Index < len(result) {
-			result[item.Index] = item.Embedding
+		if item.Index < 0 || item.Index >= len(result) {
+			return nil, fmt.Errorf(
+				"weknoracloud embedder: response index %d out of range for %d inputs",
+				item.Index,
+				len(texts),
+			)
+		}
+		if seen[item.Index] {
+			return nil, fmt.Errorf("weknoracloud embedder: duplicate response index %d", item.Index)
+		}
+		result[item.Index] = item.Embedding
+		seen[item.Index] = true
+	}
+	for index, found := range seen {
+		if !found {
+			return nil, fmt.Errorf("weknoracloud embedder: missing embedding for input index %d", index)
 		}
 	}
 	return result, nil
diff --git a/internal/models/embedding/weknoracloud_test.go b/internal/models/embedding/weknoracloud_test.go
new file mode 100644
index 000000000..fe221c4e7
--- /dev/null
+++ b/internal/models/embedding/weknoracloud_test.go
@@ -0,0 +1,100 @@
+package embedding
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func newWeKnoraCloudEmbedderTestServer(t *testing.T, response string) *WeKnoraCloudEmbedder {
+	t.Helper()
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path != weKnoraCloudEmbedPath {
+			http.Error(w, fmt.Sprintf("request path = %q, want %q", r.URL.Path, weKnoraCloudEmbedPath), http.StatusNotFound)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = fmt.Fprint(w, response)
+	}))
+	t.Cleanup(server.Close)
+
+	return &WeKnoraCloudEmbedder{
+		modelName: "test-embedding",
+		modelID:   "test-model-id",
+		appID:     "test-app-id",
+		apiKey:    "test-api-key",
+		baseURL:   server.URL,
+		client:    server.Client(),
+	}
+}
+
+func TestWeKnoraCloudBatchEmbedPreservesInputOrder(t *testing.T) {
+	embedder := newWeKnoraCloudEmbedderTestServer(t, `{
+		"data": [
+			{"index": 1, "embedding": [0.3, 0.4]},
+			{"index": 0, "embedding": [0.1, 0.2]}
+		]
+	}`)
+
+	got, err := embedder.BatchEmbed(context.Background(), []string{"first", "second"})
+	if err != nil {
+		t.Fatalf("BatchEmbed returned error: %v", err)
+	}
+	want := [][]float32{{0.1, 0.2}, {0.3, 0.4}}
+	if !reflect.DeepEqual(got, want) {
+		t.Fatalf("BatchEmbed result = %v, want %v", got, want)
+	}
+}
+
+func TestWeKnoraCloudBatchEmbedRejectsMalformedResponse(t *testing.T) {
+	tests := []struct {
+		name       string
+		response   string
+		texts      []string
+		wantErrMsg string
+	}{
+		{
+			name:       "negative index",
+			response:   `{"data": [{"index": -1, "embedding": [0.1, 0.2]}]}`,
+			texts:      []string{"first"},
+			wantErrMsg: "response index -1 out of range",
+		},
+		{
+			name:       "index above input range",
+			response:   `{"data": [{"index": 1, "embedding": [0.1, 0.2]}]}`,
+			texts:      []string{"first"},
+			wantErrMsg: "response index 1 out of range",
+		},
+		{
+			name:       "duplicate index",
+			response:   `{"data": [{"index": 0, "embedding": [0.1]}, {"index": 0, "embedding": [0.2]}]}`,
+			texts:      []string{"first"},
+			wantErrMsg: "duplicate response index 0",
+		},
+		{
+			name:       "missing result",
+			response:   `{"data": [{"index": 0, "embedding": [0.1, 0.2]}]}`,
+			texts:      []string{"first", "second"},
+			wantErrMsg: "missing embedding for input index 1",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			embedder := newWeKnoraCloudEmbedderTestServer(t, tt.response)
+
+			_, err := embedder.BatchEmbed(context.Background(), tt.texts)
+			if err == nil {
+				t.Fatalf("BatchEmbed returned nil error, want error containing %q", tt.wantErrMsg)
+			}
+			if !strings.Contains(err.Error(), tt.wantErrMsg) {
+				t.Fatalf("BatchEmbed error = %q, want it to contain %q", err, tt.wantErrMsg)
+			}
+		})
+	}
+}