From 2fe756e4be2954e8c8feb7cfd963c47087e6267e Mon Sep 17 00:00:00 2001
From: Heatherm Huang <heathermhuang@gmail.com>
Date: Wed, 1 Jul 2026 10:54:46 +0800
Subject: [PATCH] fix: recognize grok media models

---
 backend/internal/handler/grok_media.go      |  8 ++-
 backend/internal/handler/grok_media_test.go | 54 +++++++++++++++++++++
 backend/internal/pkg/xai/models.go          |  3 ++
 backend/internal/pkg/xai/oauth_test.go      |  3 ++
 4 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 backend/internal/handler/grok_media_test.go

diff --git a/backend/internal/handler/grok_media.go b/backend/internal/handler/grok_media.go
index 6122e76812..96023f452e 100644
--- a/backend/internal/handler/grok_media.go
+++ b/backend/internal/handler/grok_media.go
@@ -294,7 +294,9 @@ func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.
 		}
 
 		h.gatewayService.ReportOpenAIAccountScheduleResult(account.ID, true, nil)
-		recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, result, requestModel, body, requestID)
+		if shouldRecordGrokMediaUsage(endpoint, requestModel) {
+			recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, result, requestModel, body, requestID)
+		}
 		reqLog.Debug("grok_media.request_completed",
 			zap.Int64("account_id", account.ID),
 			zap.Int("switch_count", switchCount),
@@ -310,6 +312,10 @@ func grokMediaModerationBody(body []byte) []byte {
 	return nil
 }
 
+func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string) bool {
+	return endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) != ""
+}
+
 func recordGrokMediaUsage(
 	c *gin.Context,
 	h *OpenAIGatewayHandler,
diff --git a/backend/internal/handler/grok_media_test.go b/backend/internal/handler/grok_media_test.go
new file mode 100644
index 0000000000..1b82f8fa6f
--- /dev/null
+++ b/backend/internal/handler/grok_media_test.go
@@ -0,0 +1,54 @@
+package handler
+
+import (
+	"testing"
+
+	"github.com/Wei-Shaw/sub2api/internal/service"
+	"github.com/stretchr/testify/require"
+)
+
+func TestShouldRecordGrokMediaUsage(t *testing.T) {
+	tests := []struct {
+		name     string
+		endpoint service.GrokMediaEndpoint
+		model    string
+		want     bool
+	}{
+		{
+			name:     "image generation records usage",
+			endpoint: service.GrokMediaEndpointImagesGenerations,
+			model:    "grok-imagine",
+			want:     true,
+		},
+		{
+			name:     "image edit records usage",
+			endpoint: service.GrokMediaEndpointImagesEdits,
+			model:    "grok-imagine-edit",
+			want:     true,
+		},
+		{
+			name:     "video generation records usage",
+			endpoint: service.GrokMediaEndpointVideosGenerations,
+			model:    "grok-imagine-video-1.5",
+			want:     true,
+		},
+		{
+			name:     "video status skips empty model usage",
+			endpoint: service.GrokMediaEndpointVideoStatus,
+			model:    "",
+			want:     false,
+		},
+		{
+			name:     "generation skips usage without model",
+			endpoint: service.GrokMediaEndpointImagesGenerations,
+			model:    " ",
+			want:     false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Equal(t, tt.want, shouldRecordGrokMediaUsage(tt.endpoint, tt.model))
+		})
+	}
+}
diff --git a/backend/internal/pkg/xai/models.go b/backend/internal/pkg/xai/models.go
index 0d289274fb..74c760136d 100644
--- a/backend/internal/pkg/xai/models.go
+++ b/backend/internal/pkg/xai/models.go
@@ -15,6 +15,9 @@ var defaultModels = []Model{
 	{ID: "grok-4.20-0309-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Reasoning"},
 	{ID: "grok-4.20-0309-non-reasoning", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Non Reasoning"},
 	{ID: "grok-4.20-multi-agent-0309", Object: "model", OwnedBy: "xai", DisplayName: "Grok 4.20 Multi Agent"},
+	{ID: "grok-imagine", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine"},
+	{ID: "grok-imagine-edit", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Edit"},
+	{ID: "grok-imagine-video-1.5", Object: "model", OwnedBy: "xai", DisplayName: "Grok Imagine Video 1.5"},
 }
 
 func DefaultModels() []Model {
diff --git a/backend/internal/pkg/xai/oauth_test.go b/backend/internal/pkg/xai/oauth_test.go
index c6d601ca5d..5a728fb3ab 100644
--- a/backend/internal/pkg/xai/oauth_test.go
+++ b/backend/internal/pkg/xai/oauth_test.go
@@ -213,4 +213,7 @@ func TestDefaultModelMappingIncludesGrokAliases(t *testing.T) {
 	require.Equal(t, "grok-4.20-0309-reasoning", mapping["grok-4.20-reasoning"])
 	require.Equal(t, "grok-4.20-0309-non-reasoning", mapping["grok-4.20-non-reasoning"])
 	require.Equal(t, "grok-4.20-multi-agent-0309", mapping["grok-4.20-multi-agent-0309"])
+	require.Equal(t, "grok-imagine", mapping["grok-imagine"])
+	require.Equal(t, "grok-imagine-edit", mapping["grok-imagine-edit"])
+	require.Equal(t, "grok-imagine-video-1.5", mapping["grok-imagine-video-1.5"])
 }