zfc/sub2api
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: encode ls model credits topic values as base64

Browse Source
This commit is contained in:
win 2026-03-31 08:34:00 +08:00
parent 20151b3347
commit 8e54eaa002
2 changed files with 94 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

82
backend/internal/pkg/lspool/integration_test.go
View File

@ -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

24
backend/internal/pkg/lspool/mock_extension_server.go
View File

@ -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