From 8e54eaa0024073f6d06ab5c54fe76745962a7c81 Mon Sep 17 00:00:00 2001 From: win Date: Tue, 31 Mar 2026 08:34:00 +0800 Subject: [PATCH] fix: encode ls model credits topic values as base64 --- .../internal/pkg/lspool/integration_test.go | 82 +++++++++++++++++-- .../pkg/lspool/mock_extension_server.go | 24 ++++-- 2 files changed, 94 insertions(+), 12 deletions(-) diff --git a/backend/internal/pkg/lspool/integration_test.go b/backend/internal/pkg/lspool/integration_test.go index 4281d66c..45812bc8 100644 --- a/backend/internal/pkg/lspool/integration_test.go +++ b/backend/internal/pkg/lspool/integration_test.go @@ -3,6 +3,7 @@ package lspool import ( "bytes" "context" + "encoding/base64" "encoding/binary" "encoding/json" "fmt" @@ -72,6 +73,65 @@ func decodeProtoBytesField(data []byte, targetField int) []byte { return nil } +func decodeProtoBytesFields(data []byte, targetField int) [][]byte { + var values [][]byte + i := 0 + for i < len(data) { + tag, n := binary.Uvarint(data[i:]) + if n <= 0 { + return values + } + i += n + fieldNum := int(tag >> 3) + wireType := tag & 0x7 + switch wireType { + case 0: + _, n = binary.Uvarint(data[i:]) + if n <= 0 { + return values + } + i += n + case 2: + length, n := binary.Uvarint(data[i:]) + if n <= 0 { + return values + } + i += n + if i+int(length) > len(data) { + return values + } + if fieldNum == targetField { + values = append(values, append([]byte(nil), data[i:i+int(length)]...)) + } + i += int(length) + case 1: + i += 8 + case 5: + i += 4 + default: + return values + } + } + return values +} + +func decodeTopicRows(topic []byte) map[string]string { + rows := make(map[string]string) + for _, entry := range decodeProtoBytesFields(topic, 1) { + key := decodeProtoString(entry, 1) + row := decodeProtoBytesField(entry, 2) + rows[key] = decodeProtoString(row, 1) + } + return rows +} + +func requireBase64PrimitiveValue(t *testing.T, got string, want []byte) { + t.Helper() + decoded, err := base64.StdEncoding.DecodeString(got) + require.NoError(t, err) + require.Equal(t, want, decoded) +} + // TestMockExtensionServerTokenInjection verifies the token injection flow: // Extension → MockExtensionServer → LS subscribes uss-oauth → gets OAuthTokenInfo func TestMockExtensionServerTokenInjection(t *testing.T) { @@ -232,6 +292,11 @@ func TestUSSTopicWithModelCredits(t *testing.T) { require.Contains(t, string(topic), useAICreditsSentinelKey) require.Contains(t, string(topic), availableCreditsSentinelKey) require.Contains(t, string(topic), minimumCreditAmountForUsageKey) + + rows := decodeTopicRows(topic) + requireBase64PrimitiveValue(t, rows[useAICreditsSentinelKey], buildPrimitiveBoolBinary(true)) + requireBase64PrimitiveValue(t, rows[availableCreditsSentinelKey], buildPrimitiveInt32Binary(available)) + requireBase64PrimitiveValue(t, rows[minimumCreditAmountForUsageKey], buildPrimitiveInt32Binary(minimum)) } func TestMockExtensionServerModelCreditsDynamicUpdate(t *testing.T) { @@ -269,18 +334,23 @@ func TestMockExtensionServerModelCreditsDynamicUpdate(t *testing.T) { MinimumCreditAmountForUsage: &minimum, }) - keys := make([]string, 0, 3) - for len(keys) < 3 { + values := make(map[string]string, 3) + for len(values) < 3 { frame, readErr := readConnectFrame(resp.Body) require.NoError(t, readErr) applied := decodeProtoBytesField(frame, 2) require.NotEmpty(t, applied) - keys = append(keys, decodeProtoString(applied, 1)) + key := decodeProtoString(applied, 1) + row := decodeProtoBytesField(applied, 2) + values[key] = decodeProtoString(row, 1) } - require.Contains(t, keys, useAICreditsSentinelKey) - require.Contains(t, keys, availableCreditsSentinelKey) - require.Contains(t, keys, minimumCreditAmountForUsageKey) + require.Contains(t, values, useAICreditsSentinelKey) + require.Contains(t, values, availableCreditsSentinelKey) + require.Contains(t, values, minimumCreditAmountForUsageKey) + requireBase64PrimitiveValue(t, values[useAICreditsSentinelKey], buildPrimitiveBoolBinary(true)) + requireBase64PrimitiveValue(t, values[availableCreditsSentinelKey], buildPrimitiveInt32Binary(available)) + requireBase64PrimitiveValue(t, values[minimumCreditAmountForUsageKey], buildPrimitiveInt32Binary(minimum)) } // TestBuildInitialStateUpdate verifies the USS update wrapper diff --git a/backend/internal/pkg/lspool/mock_extension_server.go b/backend/internal/pkg/lspool/mock_extension_server.go index b5f3702a..e026be86 100644 --- a/backend/internal/pkg/lspool/mock_extension_server.go +++ b/backend/internal/pkg/lspool/mock_extension_server.go @@ -248,6 +248,18 @@ func buildPrimitiveInt32Binary(val int32) []byte { return encodeProtoVarint(3, uint64(uint32(val))) } +func encodeUSSBinaryValue(value []byte) string { + return base64.StdEncoding.EncodeToString(value) +} + +func encodeUSSPrimitiveBoolValue(val bool) string { + return encodeUSSBinaryValue(buildPrimitiveBoolBinary(val)) +} + +func encodeUSSPrimitiveInt32Value(val int32) string { + return encodeUSSBinaryValue(buildPrimitiveInt32Binary(val)) +} + func buildUSSTopicRow(key string, value string) []byte { row := buildUSSRowBinary(value) @@ -277,7 +289,7 @@ func buildUSSTopicWithModelCredits(info *ModelCreditsInfo) []byte { entries := make([][]byte, 0, 3) entries = append(entries, buildUSSTopicRow( useAICreditsSentinelKey, - string(buildPrimitiveBoolBinary(info.UseAICredits)), + encodeUSSPrimitiveBoolValue(info.UseAICredits), )) // JS protocol: useAICreditsSentinelKey carries the toggle state. // availableCreditsSentinelKey is only present when credits are enabled. @@ -286,9 +298,9 @@ func buildUSSTopicWithModelCredits(info *ModelCreditsInfo) []byte { if info.AvailableCredits != nil { credits = *info.AvailableCredits } - entries = append(entries, buildUSSTopicRow(availableCreditsSentinelKey, string(buildPrimitiveInt32Binary(credits)))) + entries = append(entries, buildUSSTopicRow(availableCreditsSentinelKey, encodeUSSPrimitiveInt32Value(credits))) } - entries = append(entries, buildUSSTopicRow(minimumCreditAmountForUsageKey, string(buildPrimitiveInt32Binary(minimum)))) + entries = append(entries, buildUSSTopicRow(minimumCreditAmountForUsageKey, encodeUSSPrimitiveInt32Value(minimum))) var topic []byte for _, entry := range entries { @@ -577,7 +589,7 @@ func buildModelCreditsAppliedUpdates(info *ModelCreditsInfo) [][]byte { updates := make([][]byte, 0, 3) updates = append(updates, buildAppliedUpdate( useAICreditsSentinelKey, - buildUSSRowBinary(string(buildPrimitiveBoolBinary(info.UseAICredits))), + buildUSSRowBinary(encodeUSSPrimitiveBoolValue(info.UseAICredits)), )) if info.UseAICredits { @@ -587,14 +599,14 @@ func buildModelCreditsAppliedUpdates(info *ModelCreditsInfo) [][]byte { } updates = append(updates, buildAppliedUpdate( availableCreditsSentinelKey, - buildUSSRowBinary(string(buildPrimitiveInt32Binary(credits))), + buildUSSRowBinary(encodeUSSPrimitiveInt32Value(credits)), )) } else { updates = append(updates, buildAppliedUpdate(availableCreditsSentinelKey, nil)) } updates = append(updates, buildAppliedUpdate( minimumCreditAmountForUsageKey, - buildUSSRowBinary(string(buildPrimitiveInt32Binary(minimum))), + buildUSSRowBinary(encodeUSSPrimitiveInt32Value(minimum)), )) return updates