From 466cee6e168900a85f3d28eb6328ee8cc6fc135d Mon Sep 17 00:00:00 2001
From: Luis Pater <webmaster@idotorg.org>
Date: Thu, 16 Jul 2026 02:32:56 +0800
Subject: [PATCH] fix: complete xAI image usage reporting

---
 internal/runtime/executor/xai_executor.go     | 10 +++--
 .../runtime/executor/xai_executor_test.go     | 41 +++++++++++++++++++
 2 files changed, 47 insertions(+), 4 deletions(-)

diff --git a/internal/runtime/executor/xai_executor.go b/internal/runtime/executor/xai_executor.go
index 2d14e82f1..edb37e73b 100644
--- a/internal/runtime/executor/xai_executor.go
+++ b/internal/runtime/executor/xai_executor.go
@@ -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
 	}
 
diff --git a/internal/runtime/executor/xai_executor_test.go b/internal/runtime/executor/xai_executor_test.go
index ee9ecbc02..aa15c29eb 100644
--- a/internal/runtime/executor/xai_executor_test.go
+++ b/internal/runtime/executor/xai_executor_test.go
@@ -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