lywsvip/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-10-17 13:08:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix(executor): preserve upstream usage on apply patch failures

Browse Source 
- Buffer and publish observed upstream usage when apply patch transformation or translation encounters an error in Meta and xAI executors
- Add stream usage buffer helpers to report failure metrics with captured upstream token usage

Closes: #6427
This commit is contained in:
Luis Pater
2026-10-07 09:31:03 +08:00
parent 39632e52a6
commit 57bde35179
6 changed files with 129 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

77
internal/runtime/executor/apply_patch_integration_test.go
View File

@@ -282,6 +282,10 @@ func (p *task6UsageCapture) HandleUsage(_ context.Context, record usage.Record)
}
}
func task6CaptureFailureUsage(t *testing.T, id string) func() {
return task6CaptureFailureUsageWithCheck(t, id, nil)
}
func task6CaptureFailureUsageWithCheck(t *testing.T, id string, check func(usage.Record)) func() {
t.Helper()
capture := &task6UsageCapture{id: id, records: make(chan usage.Record, 32)}
usage.RegisterNamedPlugin("task6-patch-failure", capture)
@@ -293,7 +297,7 @@ func task6CaptureFailureUsage(t *testing.T, id string) func() {
case record := <-capture.records:
if record.RequestID == id+"-barrier" {
if count != 1 {
t.Errorf("usage records=%d, want one failure", count)
t.Errorf("usage records=%d, want one", count)
}
return
}
@@ -301,6 +305,9 @@ func task6CaptureFailureUsage(t *testing.T, id string) func() {
if !record.Failed {
t.Errorf("invalid patch published success usage: %+v", record)
}
if check != nil {
check(record)
}
case <-time.After(3 * time.Second):
t.Fatal("usage barrier did not arrive")
}
@@ -360,6 +367,74 @@ func task6Gateway(t *testing.T, exec cliproxyauth.ProviderExecutor, auth *clipro
return recorder
}
func TestApplyPatchFailurePreservesUpstreamUsage(t *testing.T) {
const sseBody = `data: {"type":"response.created","response":{"id":"r"}}
data: {"type":"response.completed","response":{"output":[{"type":"function_call","id":"c","name":"apply_patch","arguments":"{}"}],"usage":{"input_tokens":5,"output_tokens":3}}}
`
for _, tc := range []struct {
name string
provider string
body string
stream bool
compact bool
}{
{name: "meta-direct", provider: "meta", body: `{"object":"response","output":[{"type":"function_call","id":"c","name":"apply_patch","arguments":"{}"}],"usage":{"input_tokens":5,"output_tokens":3}}`},
{name: "meta-direct-event", provider: "meta", body: `{"type":"response.completed","response":{"output":[{"type":"function_call","id":"c","name":"apply_patch","arguments":"{}"}],"usage":{"input_tokens":5,"output_tokens":3}}}`},
{name: "xai-compact", provider: "xai", body: `{"id":"r","output":[{"type":"function_call","id":"c","name":"apply_patch","arguments":"{}"}],"usage":{"input_tokens":5,"output_tokens":3}}`, compact: true},
{name: "meta-sse", provider: "meta", body: sseBody, stream: true},
{name: "xai-sse", provider: "xai", body: sseBody, stream: true},
{name: "meta-buffered-sse", provider: "meta", body: sseBody},
{name: "xai-buffered-sse", provider: "xai", body: sseBody},
} {
t.Run(tc.name, func(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if tc.compact && r.URL.Path != "/responses/compact" {
t.Errorf("request path = %q, want /responses/compact", r.URL.Path)
}
if tc.stream {
w.Header().Set("Content-Type", "text/event-stream")
}
_, _ = io.WriteString(w, tc.body)
}))
defer server.Close()
var exec cliproxyauth.ProviderExecutor
if tc.provider == "meta" {
exec = NewMetaExecutor(&config.Config{})
} else {
exec = NewXAIExecutor(&config.Config{})
}
auth := &cliproxyauth.Auth{ID: "task6-usage-" + tc.name, Provider: tc.provider, Attributes: map[string]string{"api_key": "test", "base_url": server.URL}}
checkUsage := task6CaptureFailureUsageWithCheck(t, auth.ID, func(record usage.Record) {
if record.Detail.InputTokens != 5 || record.Detail.OutputTokens != 3 {
t.Errorf("failure usage = %+v, want input/output 5/3", record.Detail)
}
})
req := cliproxyexecutor.Request{Model: "grok-4", Payload: []byte(task6PatchRequest)}
opts := cliproxyexecutor.Options{SourceFormat: sdktranslator.FormatOpenAIResponse, OriginalRequest: req.Payload}
if tc.compact {
opts.Alt = "responses/compact"
}
if tc.stream {
result, errExecuteStream := exec.ExecuteStream(t.Context(), auth, req, opts)
if errExecuteStream != nil {
t.Fatal(errExecuteStream)
}
assertTask6FailedStream(t, result.Chunks)
} else {
result, errExecute := exec.Execute(t.Context(), auth, req, opts)
if len(result.Payload) != 0 {
t.Errorf("failed response returned payload %s", result.Payload)
}
assertTask6PatchError(t, errExecute)
}
checkUsage()
})
}
}
func TestApplyPatchHTTPGatewayErrorMatrix(t *testing.T) {
for _, provider := range []string{"custom-compat", "claude", "claude-oauth", "gemini", "gemini-interactions", "vertex", "antigravity", "devin", "xai", "meta", "kimi"} {
for _, mode := range []string{"nonstream", "stream", "eof"} {

15
internal/runtime/executor/helps/apply_patch.go
View File

@@ -38,9 +38,17 @@ func CanFinalizeResponseStream(param any) bool {
// RecordApplyPatchStreamFailure records validation before delivery can be canceled.
func RecordApplyPatchStreamFailure(ctx context.Context, param any, reporter *UsageReporter, gatewayErr error) bool {
return RecordApplyPatchStreamFailureWithUsage(ctx, param, reporter, nil, gatewayErr)
}
// RecordApplyPatchStreamFailureWithUsage records validation with buffered upstream usage.
func RecordApplyPatchStreamFailureWithUsage(ctx context.Context, param any, reporter *UsageReporter, streamUsage *StreamUsageBuffer, gatewayErr error) bool {
if ApplyPatchTranslationError(param) == nil {
return false
}
if streamUsage != nil && streamUsage.PublishFailure(ctx, reporter, gatewayErr) {
return true
}
reporter.PublishFailure(ctx, gatewayErr)
return true
}
@@ -48,7 +56,12 @@ func RecordApplyPatchStreamFailure(ctx context.Context, param any, reporter *Usa
// StopApplyPatchStream propagates a retained failure after its one translated frame.
// The caller supplies its existing sanitized gateway status error.
func StopApplyPatchStream(ctx context.Context, param any, reporter *UsageReporter, out chan<- cliproxyexecutor.StreamChunk, gatewayErr error) bool {
if !RecordApplyPatchStreamFailure(ctx, param, reporter, gatewayErr) {
return StopApplyPatchStreamWithUsage(ctx, param, reporter, nil, out, gatewayErr)
}
// StopApplyPatchStreamWithUsage propagates a retained failure with buffered upstream usage.
func StopApplyPatchStreamWithUsage(ctx context.Context, param any, reporter *UsageReporter, streamUsage *StreamUsageBuffer, out chan<- cliproxyexecutor.StreamChunk, gatewayErr error) bool {
if !RecordApplyPatchStreamFailureWithUsage(ctx, param, reporter, streamUsage, gatewayErr) {
return false
}
select {

12
internal/runtime/executor/meta_executor_execute.go
View File

@@ -143,8 +143,10 @@ func (e *MetaExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
return resp, wrapMetaUpstreamError(httpResp.StatusCode, data)
}
out, errCompleted := e.translateMetaCompleted(ctx, req, prepared, data)
var upstreamUsage helps.StreamUsageBuffer
out, errCompleted := e.translateMetaCompleted(ctx, req, prepared, data, &upstreamUsage)
if errCompleted != nil {
upstreamUsage.PublishFailure(ctx, reporter, errCompleted)
return resp, errCompleted
}
if len(out.sourceEvent) > 0 {
@@ -167,7 +169,7 @@ type metaCompletedTranslation struct {
sourceEvent []byte
}
func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxyexecutor.Request, prepared *metaPreparedRequest, data []byte) (metaCompletedTranslation, error) {
func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxyexecutor.Request, prepared *metaPreparedRequest, data []byte, upstreamUsage *helps.StreamUsageBuffer) (metaCompletedTranslation, error) {
outputItemsByIndex := make(map[int64][]byte)
var outputItemsFallback [][]byte
for _, line := range bytes.Split(data, []byte("\n")) {
@@ -175,6 +177,9 @@ func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxye
continue
}
eventData := bytes.TrimSpace(line[len(dataTag):])
if detail, ok := helps.ParseCodexUsage(eventData); ok {
upstreamUsage.Observe(detail, true)
}
if errEvent := metaStreamEventError(eventData); errEvent != nil {
return metaCompletedTranslation{}, errEvent
}
@@ -201,6 +206,9 @@ func (e *MetaExecutor) translateMetaCompleted(ctx context.Context, req cliproxye
}
if completedData, ok := metaAsCompletedEvent(data); ok {
if detail, okUsage := helps.ParseCodexUsage(completedData); okUsage {
upstreamUsage.Observe(detail, true)
}
completedData = patchCodexCompletedOutput(completedData, outputItemsByIndex, outputItemsFallback)
var errBridge error
completedData, errBridge = prepared.applyPatch.Bridge.TransformNonStream(completedData)

12
internal/runtime/executor/meta_executor_stream.go
View File

@@ -89,13 +89,13 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
emitTranslatedLine := func(translatedLine []byte) bool {
lines, errBridge := prepared.applyPatch.Stream(translatedLine)
if errBridge != nil {
reporter.PublishFailure(ctx, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
streamUsage.PublishFailure(ctx, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
}
var chunks [][]byte
for _, line := range lines {
chunks = append(chunks, helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, &param, claudeInputTokens)...)
}
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
helps.RecordApplyPatchStreamFailureWithUsage(ctx, param, reporter, &streamUsage, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
for i := range chunks {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:
@@ -103,12 +103,12 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
return false
}
}
if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
if helps.StopApplyPatchStreamWithUsage(ctx, param, reporter, &streamUsage, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
return false
}
if errBridge != nil {
errBridge = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailure(ctx, errBridge)
streamUsage.PublishFailure(ctx, reporter, errBridge)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errBridge}:
case <-ctx.Done():
@@ -153,7 +153,7 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
}
finishEvents, errFinish := prepared.applyPatch.FinishStream()
if errFinish != nil {
reporter.PublishFailure(ctx, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
streamUsage.PublishFailure(ctx, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
}
for _, event := range finishEvents {
for _, chunk := range helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, event, &param, claudeInputTokens) {
@@ -166,7 +166,7 @@ func (e *MetaExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
}
if errFinish != nil {
errFinish = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailure(ctx, errFinish)
streamUsage.PublishFailure(ctx, reporter, errFinish)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errFinish}:
case <-ctx.Done():

21
internal/runtime/executor/xai_executor_execute.go
View File

@@ -87,11 +87,15 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
var outputItemsFallback [][]byte
responseFilter := newXAIInternalXSearchResponseFilter(prepared.filterInternalXSearch, prepared.clientDeclaredTools)
namespaceRestorer := newXAINamespaceRestorer(prepared.namespaceTools)
var upstreamUsage helps.StreamUsageBuffer
for _, line := range bytes.Split(data, []byte("\n")) {
if !bytes.HasPrefix(line, xaiDataTag) {
continue
}
eventData := xaiNormalizeReasoningSummaryData(bytes.TrimSpace(line[len(xaiDataTag):]))
if detail, ok := helps.ParseCodexUsage(eventData); ok {
upstreamUsage.Observe(detail, true)
}
prepared.applyPatch.RememberDispatcherEvent(eventData)
eventData = namespaceRestorer.restore(eventData)
if prepared.webSearchAlias != "" {
@@ -104,6 +108,7 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
events, errBridge := prepared.applyPatch.Transform(eventData)
if errBridge != nil {
errBridge = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
upstreamUsage.PublishFailure(ctx, reporter, errBridge)
return resp, errBridge
}
for _, eventData := range events {
@@ -123,7 +128,9 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
var param any
out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, completedData, &param)
if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 {
return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
errTranslation := statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
upstreamUsage.PublishFailure(ctx, reporter, errTranslation)
return cliproxyexecutor.Response{}, errTranslation
}
if detail, ok := helps.ParseCodexUsage(eventData); ok {
reporter.Publish(ctx, detail)
@@ -137,7 +144,9 @@ func (e *XAIExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
}
if errFinish := prepared.applyPatch.Finish(); errFinish != nil {
return resp, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
errFinish = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
upstreamUsage.PublishFailure(ctx, reporter, errFinish)
return resp, errFinish
}
return resp, statusErr{code: http.StatusRequestTimeout, msg: "xai stream error: stream disconnected before response.completed or response.incomplete"}
}
@@ -149,17 +158,21 @@ func (e *XAIExecutor) executeCompact(ctx context.Context, auth *cliproxyauth.Aut
}
defer reporter.TrackFailure(ctx, &err)
upstreamUsage := helps.ParseOpenAIUsage(data)
converted, errBridge := prepared.applyPatch.Bridge.TransformNonStream(data)
if errBridge != nil {
errBridge = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailureWithDetail(ctx, upstreamUsage, errBridge)
return resp, errBridge
}
var param any
out := sdktranslator.TranslateNonStream(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, converted, &param)
if helps.ApplyPatchTranslationError(param) != nil || len(out) == 0 {
return cliproxyexecutor.Response{}, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
errTranslation := statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailureWithDetail(ctx, upstreamUsage, errTranslation)
return cliproxyexecutor.Response{}, errTranslation
}
reporter.Publish(ctx, helps.ParseOpenAIUsage(data))
reporter.Publish(ctx, upstreamUsage)
if prepared.responseFormat == sdktranslator.FormatOpenAIResponse {
out = helps.EnsureResponsesUsageDetails(out)
}

12
internal/runtime/executor/xai_executor_stream.go
View File

@@ -90,7 +90,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
emitTranslatedLine := func(translatedLine []byte) bool {
lines, errBridge := prepared.applyPatch.Stream(translatedLine)
if errBridge != nil {
reporter.PublishFailure(ctx, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
streamUsage.PublishFailure(ctx, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
}
var chunks [][]byte
for _, line := range lines {
@@ -113,7 +113,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
}
chunks = append(chunks, helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, line, &param, claudeInputTokens)...)
}
helps.RecordApplyPatchStreamFailure(ctx, param, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
helps.RecordApplyPatchStreamFailureWithUsage(ctx, param, reporter, &streamUsage, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
for i := range chunks {
select {
case out <- cliproxyexecutor.StreamChunk{Payload: chunks[i]}:
@@ -121,12 +121,12 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
return false
}
}
if helps.StopApplyPatchStream(ctx, param, reporter, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
if helps.StopApplyPatchStreamWithUsage(ctx, param, reporter, &streamUsage, out, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}) {
return false
}
if errBridge != nil {
errBridge = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailure(ctx, errBridge)
streamUsage.PublishFailure(ctx, reporter, errBridge)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errBridge}:
case <-ctx.Done():
@@ -203,7 +203,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
}
finishEvents, errFinish := prepared.applyPatch.FinishStream()
if errFinish != nil {
reporter.PublishFailure(ctx, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
streamUsage.PublishFailure(ctx, reporter, statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage})
}
for _, event := range finishEvents {
for _, chunk := range helps.TranslateStreamWithClaudeInputTokens(ctx, prepared.to, prepared.responseFormat, req.Model, prepared.originalPayload, prepared.body, event, &param, claudeInputTokens) {
@@ -216,7 +216,7 @@ func (e *XAIExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Auth
}
if errFinish != nil {
errFinish = statusErr{code: http.StatusBadGateway, msg: helps.ApplyPatchUpstreamErrorMessage}
reporter.PublishFailure(ctx, errFinish)
streamUsage.PublishFailure(ctx, reporter, errFinish)
select {
case out <- cliproxyexecutor.StreamChunk{Err: errFinish}:
case <-ctx.Done():