From b3207432ae3aacb9fb281354c8bdf021ead5d9d7 Mon Sep 17 00:00:00 2001 From: Dev Agent Date: Wed, 8 Apr 2026 12:46:48 +0800 Subject: [PATCH 1/2] fix(aigateway): fix empty resource id for internal models when cache matched --- aigateway/component/openai_ce_test.go | 36 +++++--- aigateway/types/openai.go | 18 ++++ aigateway/types/openai_test.go | 120 ++++++++++++++++++++++++-- 3 files changed, 155 insertions(+), 19 deletions(-) diff --git a/aigateway/component/openai_ce_test.go b/aigateway/component/openai_ce_test.go index 11d62e8c8..11adef648 100644 --- a/aigateway/component/openai_ce_test.go +++ b/aigateway/component/openai_ce_test.go @@ -104,9 +104,12 @@ func TestOpenAIComponent_GetAvailableModels(t *testing.T) { }, Endpoint: "endpoint1", InternalModelInfo: types.InternalModelInfo{ - ClusterID: deploys[0].ClusterID, - SvcName: deploys[0].SvcName, - ImageID: deploys[0].ImageID, + CSGHubModelID: deploys[0].Repository.Path, + OwnerUUID: deploys[0].User.UUID, + ClusterID: deploys[0].ClusterID, + SvcName: deploys[0].SvcName, + SvcType: deploys[0].Type, + ImageID: deploys[0].ImageID, }, InternalUse: true, }, @@ -121,9 +124,12 @@ func TestOpenAIComponent_GetAvailableModels(t *testing.T) { }, Endpoint: "endpoint2", InternalModelInfo: types.InternalModelInfo{ - ClusterID: deploys[1].ClusterID, - SvcName: deploys[1].SvcName, - ImageID: deploys[1].ImageID, + CSGHubModelID: deploys[1].Repository.Path, + OwnerUUID: deploys[1].User.UUID, + ClusterID: deploys[1].ClusterID, + SvcName: deploys[1].SvcName, + SvcType: deploys[1].Type, + ImageID: deploys[1].ImageID, }, InternalUse: true, }, @@ -207,9 +213,12 @@ func TestOpenAIComponent_GetAvailableModels(t *testing.T) { }, Endpoint: "endpoint3", InternalModelInfo: types.InternalModelInfo{ - ClusterID: deploys[0].ClusterID, - SvcName: deploys[0].SvcName, - ImageID: deploys[0].ImageID, + CSGHubModelID: deploys[0].Repository.Path, + OwnerUUID: deploys[0].User.UUID, + ClusterID: deploys[0].ClusterID, + SvcName: deploys[0].SvcName, + SvcType: deploys[0].Type, + ImageID: deploys[0].ImageID, }, InternalUse: true, }, @@ -293,9 +302,12 @@ func TestOpenAIComponent_GetModelByID(t *testing.T) { }, Endpoint: "endpoint1", InternalModelInfo: types.InternalModelInfo{ - ClusterID: deploys[0].ClusterID, - SvcName: deploys[0].SvcName, - ImageID: deploys[0].ImageID, + CSGHubModelID: deploys[0].Repository.Path, + OwnerUUID: deploys[0].User.UUID, + ClusterID: deploys[0].ClusterID, + SvcName: deploys[0].SvcName, + SvcType: deploys[0].Type, + ImageID: deploys[0].ImageID, }, InternalUse: true, }, diff --git a/aigateway/types/openai.go b/aigateway/types/openai.go index 93426aa6e..929b0b88a 100644 --- a/aigateway/types/openai.go +++ b/aigateway/types/openai.go @@ -60,8 +60,11 @@ func (m Model) MarshalJSON() ([]byte, error) { Public bool `json:"public"` Endpoint string `json:"endpoint"` Metadata map[string]any `json:"metadata"` + CSGHubModelID *string `json:"csghub_model_id,omitempty"` + OwnerUUID *string `json:"owner_uuid,omitempty"` ClusterID *string `json:"cluster_id,omitempty"` SvcName *string `json:"svc_name,omitempty"` + SvcType *int `json:"svc_type,omitempty"` ImageID *string `json:"image_id,omitempty"` AuthHead *string `json:"auth_head,omitempty"` Provider *string `json:"provider,omitempty"` @@ -82,6 +85,12 @@ func (m Model) MarshalJSON() ([]byte, error) { supportFC := m.SupportFunctionCall resp.SupportFunctionCall = &supportFC } + if m.CSGHubModelID != "" { + resp.CSGHubModelID = &m.CSGHubModelID + } + if m.OwnerUUID != "" { + resp.OwnerUUID = &m.OwnerUUID + } if m.Provider != "" { resp.Provider = &m.Provider } @@ -94,6 +103,9 @@ func (m Model) MarshalJSON() ([]byte, error) { if m.SvcName != "" { resp.SvcName = &m.SvcName } + if m.SvcType != 0 { + resp.SvcType = &m.SvcType + } if m.ImageID != "" { resp.ImageID = &m.ImageID } @@ -116,8 +128,11 @@ func (m *Model) UnmarshalJSON(data []byte) error { Public bool `json:"public"` Endpoint string `json:"endpoint"` Metadata map[string]any `json:"metadata"` + CSGHubModelID string `json:"csghub_model_id,omitempty"` + OwnerUUID string `json:"owner_uuid,omitempty"` ClusterID string `json:"cluster_id,omitempty"` SvcName string `json:"svc_name,omitempty"` + SvcType int `json:"svc_type,omitempty"` ImageID string `json:"image_id,omitempty"` AuthHead string `json:"auth_head,omitempty"` Provider string `json:"provider,omitempty"` @@ -136,8 +151,11 @@ func (m *Model) UnmarshalJSON(data []byte) error { m.Public = aux.Public m.Endpoint = aux.Endpoint m.Metadata = aux.Metadata + m.CSGHubModelID = aux.CSGHubModelID + m.OwnerUUID = aux.OwnerUUID m.ClusterID = aux.ClusterID m.SvcName = aux.SvcName + m.SvcType = aux.SvcType m.ImageID = aux.ImageID m.AuthHead = aux.AuthHead m.Provider = aux.Provider diff --git a/aigateway/types/openai_test.go b/aigateway/types/openai_test.go index 3ceba80b2..8450eba08 100644 --- a/aigateway/types/openai_test.go +++ b/aigateway/types/openai_test.go @@ -4,6 +4,8 @@ import ( "encoding/json" "strings" "testing" + + "github.com/stretchr/testify/require" ) // TestModelSerialization tests the custom serialization of Model struct @@ -17,10 +19,10 @@ func TestModelSerialization(t *testing.T) { Task: "text-generation", SupportFunctionCall: true, - Public: true, }, InternalModelInfo: InternalModelInfo{ CSGHubModelID: "test/repo/path", + OwnerUUID: "test-owner-uuid", ClusterID: "test-cluster-id", SvcName: "test-service", SvcType: 1, @@ -47,8 +49,8 @@ func TestModelSerialization(t *testing.T) { t.Errorf("External response should not contain sensitive fields, got: %s", jsonStr) } - if !contains(jsonStr, "test-model") || !contains(jsonStr, "model") || !contains(jsonStr, "test-owner") || !contains(jsonStr, "public") { - t.Errorf("External response should contain BaseModel fields including public, got: %s", jsonStr) + if !contains(jsonStr, "test-model") || !contains(jsonStr, "model") || !contains(jsonStr, "test-owner") { + t.Errorf("External response should contain BaseModel fields, got: %s", jsonStr) } }) @@ -60,8 +62,8 @@ func TestModelSerialization(t *testing.T) { t.Fatalf("Failed to marshal model in internal use mode: %v", err) } jsonStr := string(jsonData) - if !contains(jsonStr, "endpoint") || !contains(jsonStr, "http://test-endpoint.com") || !contains(jsonStr, "test-model") || !contains(jsonStr, "public") { - t.Errorf("Internal response should contain base fields including public, got: %s", jsonStr) + if !contains(jsonStr, "endpoint") || !contains(jsonStr, "http://test-endpoint.com") || !contains(jsonStr, "test-model") { + t.Errorf("Internal response should contain base fields, got: %s", jsonStr) } if contains(jsonStr, "internal_model_info") { @@ -79,6 +81,17 @@ func TestModelSerialization(t *testing.T) { if !contains(jsonStr, "image_id") || !contains(jsonStr, "test-image-id") { t.Errorf("Internal response should contain expanded InternalModelInfo fields, got: %s", jsonStr) } + // csghub_model_id, owner_uuid, and svc_type must survive the Redis round-trip so + // that RecordUsage can populate resource_id for inference models. + if !contains(jsonStr, "csghub_model_id") || !contains(jsonStr, "test/repo/path") { + t.Errorf("Internal response should contain csghub_model_id, got: %s", jsonStr) + } + if !contains(jsonStr, "owner_uuid") { + t.Errorf("Internal response should contain owner_uuid, got: %s", jsonStr) + } + if !contains(jsonStr, "svc_type") { + t.Errorf("Internal response should contain svc_type, got: %s", jsonStr) + } }) // case3: mode switching @@ -108,7 +121,6 @@ func TestModelListSerialization(t *testing.T) { BaseModel: BaseModel{ ID: "model-1", Object: "model", - Public: true, }, Endpoint: "http://model-1.com", InternalUse: false, @@ -151,7 +163,6 @@ func TestModelUnmarshal(t *testing.T) { "owned_by": "test-owner", "task": "text-generation", "support_function_call": true, - "public": true, "endpoint": "http://model-1.com", "internal_use": false } @@ -169,3 +180,98 @@ func TestModelUnmarshal(t *testing.T) { t.Errorf("Model list unmarshal failed, got: %v", modelList) } } + +// TestInferenceModelRoundTrip verifies that CSGHubModelID, OwnerUUID, and SvcType +// survive a Redis marshal→unmarshal cycle so that RecordUsage can always populate +// resource_id for inference (llm_type=inference) models. +func TestInferenceModelRoundTrip(t *testing.T) { + original := &Model{ + BaseModel: BaseModel{ + ID: "Qwen/Qwen3Guard-Gen-0.6B:fgufi9nytc00", + Object: "model", + Created: 1633046400, + OwnedBy: "Qwen", + }, + InternalModelInfo: InternalModelInfo{ + CSGHubModelID: "Qwen/Qwen3Guard-Gen-0.6B", + OwnerUUID: "uuid-owner-123", + ClusterID: "cluster-abc", + SvcName: "fgufi9nytc00", + SvcType: 2, + ImageID: "img-xyz", + }, + Endpoint: "http://inference.internal/v1", + InternalUse: true, + } + + data, err := json.Marshal(original) + require.NoError(t, err, "marshal should not error") + + var restored Model + require.NoError(t, json.Unmarshal(data, &restored), "unmarshal should not error") + + require.Equal(t, original.CSGHubModelID, restored.CSGHubModelID, "CSGHubModelID must round-trip") + require.Equal(t, original.OwnerUUID, restored.OwnerUUID, "OwnerUUID must round-trip") + require.Equal(t, original.SvcType, restored.SvcType, "SvcType must round-trip") + require.Equal(t, original.ClusterID, restored.ClusterID, "ClusterID must round-trip") + require.Equal(t, original.SvcName, restored.SvcName, "SvcName must round-trip") + require.Equal(t, original.ImageID, restored.ImageID, "ImageID must round-trip") + require.Equal(t, original.ID, restored.ID, "ID must round-trip") + require.Equal(t, original.Endpoint, restored.Endpoint, "Endpoint must round-trip") +} + +func TestModel_SkipBalance(t *testing.T) { + tests := []struct { + name string + metadata map[string]any + expected bool + }{ + { + name: "Metadata is nil", + metadata: nil, + expected: false, + }, + { + name: "Metadata does not have MetaTaskKey", + metadata: map[string]any{}, + expected: false, + }, + { + name: "MetaTaskKey value is not a slice", + metadata: map[string]any{MetaTaskKey: "not a slice"}, + expected: false, + }, + { + name: "MetaTaskKey value is slice but not of strings", + metadata: map[string]any{MetaTaskKey: []int{1, 2, 3}}, + expected: false, + }, + { + name: "MetaTaskKey value is slice of strings but does not contain MetaTaskValGuard", + metadata: map[string]any{MetaTaskKey: []interface{}{"text-generation", "text-to-image"}}, + expected: false, + }, + { + name: "MetaTaskKey value is slice of strings and contains MetaTaskValGuard", + metadata: map[string]any{MetaTaskKey: []interface{}{"text-generation", MetaTaskValGuard}}, + expected: true, + }, + { + name: "MetaTaskKey value is slice of mixed types with MetaTaskValGuard", + metadata: map[string]any{MetaTaskKey: []interface{}{1, "text-generation", MetaTaskValGuard, 3.14}}, + expected: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + model := &Model{ + BaseModel: BaseModel{ + Metadata: tt.metadata, + }, + } + result := model.SkipBalance() + require.Equal(t, tt.expected, result) + }) + } +} From 6d1346e2de27b04c43ef3080648348c00a310de2 Mon Sep 17 00:00:00 2001 From: Lei Da Date: Thu, 9 Apr 2026 19:18:45 +0800 Subject: [PATCH 2/2] fix test --- aigateway/types/openai_test.go | 2 -- 1 file changed, 2 deletions(-) diff --git a/aigateway/types/openai_test.go b/aigateway/types/openai_test.go index 530b40947..2f60beec8 100644 --- a/aigateway/types/openai_test.go +++ b/aigateway/types/openai_test.go @@ -5,8 +5,6 @@ import ( "github.com/stretchr/testify/require" "strings" "testing" - - "github.com/stretchr/testify/require" ) // TestModelSerialization tests the custom serialization of Model struct