From 57bde35179ecbdca176ca8923d39cc18805774f2 Mon Sep 17 00:00:00 2001 From: Luis Pater Date: Wed, 7 Oct 2026 09:31:03 +0800 Subject: [PATCH] fix(executor): preserve upstream usage on apply patch failures - 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 --- .../executor/apply_patch_integration_test.go | 77 ++++++++++++++++++- .../runtime/executor/helps/apply_patch.go | 15 +++- .../runtime/executor/meta_executor_execute.go | 12 ++- .../runtime/executor/meta_executor_stream.go | 12 +-- .../runtime/executor/xai_executor_execute.go | 21 ++++- .../runtime/executor/xai_executor_stream.go | 12 +-- 6 files changed, 129 insertions(+), 20 deletions(-) diff --git a/internal/runtime/executor/apply_patch_integration_test.go b/internal/runtime/executor/apply_patch_integration_test.go index 4140bb3a3..10a4ee823 100644 --- a/internal/runtime/executor/apply_patch_integration_test.go +++ b/internal/runtime/executor/apply_patch_integration_test.go @@ -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"} { diff --git a/internal/runtime/executor/helps/apply_patch.go b/internal/runtime/executor/helps/apply_patch.go index bfbb78901..a0223f153 100644 --- a/internal/runtime/executor/helps/apply_patch.go +++ b/internal/runtime/executor/helps/apply_patch.go @@ -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 { diff --git a/internal/runtime/executor/meta_executor_execute.go b/internal/runtime/executor/meta_executor_execute.go index 6de21936a..3fa0c19a5 100644 --- a/internal/runtime/executor/meta_executor_execute.go +++ b/internal/runtime/executor/meta_executor_execute.go @@ -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) diff --git a/internal/runtime/executor/meta_executor_stream.go b/internal/runtime/executor/meta_executor_stream.go index 8d5f54626..d67330908 100644 --- a/internal/runtime/executor/meta_executor_stream.go +++ b/internal/runtime/executor/meta_executor_stream.go @@ -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, ¶m, 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, ¶m, 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(): diff --git a/internal/runtime/executor/xai_executor_execute.go b/internal/runtime/executor/xai_executor_execute.go index de001c961..ef4e80a19 100644 --- a/internal/runtime/executor/xai_executor_execute.go +++ b/internal/runtime/executor/xai_executor_execute.go @@ -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, ¶m) 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, ¶m) 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) } diff --git a/internal/runtime/executor/xai_executor_stream.go b/internal/runtime/executor/xai_executor_stream.go index d7918e1ee..1c5588a8b 100644 --- a/internal/runtime/executor/xai_executor_stream.go +++ b/internal/runtime/executor/xai_executor_stream.go @@ -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, ¶m, 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, ¶m, 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():