fix: complete xAI image usage reporting

This commit is contained in:
Luis Pater
2026-07-16 02:32:56 +08:00
parent 3cb2d27d98
commit 466cee6e16
2 changed files with 47 additions and 4 deletions

View File

@@ -481,8 +481,12 @@ func xaiBuildSSEFrame(eventName string, data []byte) []byte {
}
func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, endpointPath string) (resp cliproxyexecutor.Response, err error) {
model := gjson.GetBytes(req.Payload, "model").String()
model := strings.TrimSpace(gjson.GetBytes(req.Payload, "model").String())
if model == "" {
model = strings.TrimSpace(req.Model)
}
reporter := helps.NewExecutorUsageReporter(ctx, e, model, auth)
defer reporter.TrackFailure(ctx, &err)
token, baseURL := xaiCreds(auth)
if baseURL == "" {
@@ -502,10 +506,10 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
e.recordXAIRequest(ctx, auth, url, httpReq.Header.Clone(), req.Payload)
httpClient := helps.NewProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
httpClient = reporter.TrackHTTPClient(httpClient)
httpResp, err := httpClient.Do(httpReq)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
reporter.PublishFailure(ctx, err)
return resp, err
}
defer func() {
@@ -518,7 +522,6 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
data, err := io.ReadAll(httpResp.Body)
if err != nil {
helps.RecordAPIResponseError(ctx, e.cfg, err)
reporter.PublishFailure(ctx, err)
return resp, err
}
helps.AppendAPIResponseChunk(ctx, e.cfg, data)
@@ -526,7 +529,6 @@ func (e *XAIExecutor) executeImages(ctx context.Context, auth *cliproxyauth.Auth
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
helps.LogWithRequestID(ctx).Debugf("request error, error status: %d, error message: %s", httpResp.StatusCode, helps.SummarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
err = xaiStatusErr(httpResp.StatusCode, data)
reporter.PublishFailure(ctx, err)
return resp, err
}

View File

@@ -2088,6 +2088,9 @@ func TestXAIExecutorExecuteImagesUsesImagesEndpointAndPublishesUsage(t *testing.
if record.Detail != (usage.Detail{}) {
t.Fatalf("detail = %+v, want zero token usage", record.Detail)
}
if record.TTFT <= 0 {
t.Fatalf("ttft = %v, want positive duration", record.TTFT)
}
assertNoAdditionalXAIUsageRecord(t, plugin.records)
}
@@ -2140,6 +2143,44 @@ func TestXAIExecutorExecuteImagesPublishesFailureUsage(t *testing.T) {
assertNoAdditionalXAIUsageRecord(t, plugin.records)
}
func TestXAIExecutorExecuteImagesPublishesRequestBuildFailureUsage(t *testing.T) {
const requestedModel = "grok-imagine-image-fallback"
plugin := &captureXAIUsagePlugin{
model: requestedModel,
records: make(chan usage.Record, 2),
}
usage.RegisterPlugin(plugin)
exec := NewXAIExecutor(&config.Config{})
auth := &cliproxyauth.Auth{
Provider: "xai",
Attributes: map[string]string{"base_url": "://invalid"},
}
_, err := exec.Execute(context.Background(), auth, cliproxyexecutor.Request{
Model: requestedModel,
Payload: []byte(`{"prompt":"draw"}`),
}, cliproxyexecutor.Options{
SourceFormat: sdktranslator.FromString("openai-image"),
Metadata: map[string]any{
cliproxyexecutor.RequestPathMetadataKey: "/v1/images/generations",
},
})
if err == nil {
t.Fatal("Execute() error = nil, want non-nil")
}
record := waitForXAIUsageRecord(t, plugin.records)
if record.Model != requestedModel {
t.Fatalf("model = %q, want %q", record.Model, requestedModel)
}
if !record.Failed {
t.Fatal("failed = false, want true")
}
assertNoAdditionalXAIUsageRecord(t, plugin.records)
}
type captureXAIUsagePlugin struct {
model string
records chan usage.Record