chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-03 13:00:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(executor): expire stop chunks without usage metadata

Browse Source
This commit is contained in:
hkfires
2025-11-22 15:27:47 +08:00
parent 19a048879c
commit 8356b35320
1 changed files with 29 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
internal/runtime/executor/usage_helpers.go
View File

@@ -416,28 +416,13 @@ func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {
return detail, true
}
func jsonPayload(line []byte) []byte {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 {
return nil
}
if bytes.Equal(trimmed, []byte("[DONE]")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("event:")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("data:")) {
trimmed = bytes.TrimSpace(trimmed[len("data:"):])
}
if len(trimmed) == 0 || trimmed[0] != '{' {
return nil
}
return trimmed
}
var stopChunkWithoutUsage sync.Map
func rememberStopWithoutUsage(traceID string) {
stopChunkWithoutUsage.Store(traceID, struct{}{})
time.AfterFunc(10*time.Minute, func() { stopChunkWithoutUsage.Delete(traceID) })
}
// FilterSSEUsageMetadata removes usageMetadata from SSE events that are not
// terminal (finishReason != "stop"). Stop chunks are left untouched. This
// function is shared between aistudio and antigravity executors.
@@ -462,15 +447,13 @@ func FilterSSEUsageMetadata(payload []byte) []byte {
rawJSON := bytes.TrimSpace(line[dataIdx+5:])
traceID := gjson.GetBytes(rawJSON, "traceId").String()
if isStopChunkWithoutUsage(rawJSON) && traceID != "" {
stopChunkWithoutUsage.Store(traceID, true)
rememberStopWithoutUsage(traceID)
continue
}
if traceID != "" {
if v, ok := stopChunkWithoutUsage.Load(traceID); ok {
if keep, _ := v.(bool); keep && hasUsageMetadata(rawJSON) {
stopChunkWithoutUsage.Delete(traceID)
continue
}
if _, ok := stopChunkWithoutUsage.Load(traceID); ok && hasUsageMetadata(rawJSON) {
stopChunkWithoutUsage.Delete(traceID)
continue
}
}
@@ -578,3 +561,23 @@ func isStopChunkWithoutUsage(jsonBytes []byte) bool {
}
return !hasUsageMetadata(jsonBytes)
}
func jsonPayload(line []byte) []byte {
trimmed := bytes.TrimSpace(line)
if len(trimmed) == 0 {
return nil
}
if bytes.Equal(trimmed, []byte("[DONE]")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("event:")) {
return nil
}
if bytes.HasPrefix(trimmed, []byte("data:")) {
trimmed = bytes.TrimSpace(trimmed[len("data:"):])
}
if len(trimmed) == 0 || trimmed[0] != '{' {
return nil
}
return trimmed
}