diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index cbf43515..7fd98855 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -167,6 +167,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se group.FieldSupportedModelScopes, group.FieldAllowMessagesDispatch, group.FieldDefaultMappedModel, + group.FieldMessagesDispatchModelConfig, ) }). Only(ctx) @@ -689,6 +690,7 @@ func groupEntityToService(g *dbent.Group) *service.Group { RequireOAuthOnly: g.RequireOauthOnly, RequirePrivacySet: g.RequirePrivacySet, DefaultMappedModel: g.DefaultMappedModel, + MessagesDispatchModelConfig: g.MessagesDispatchModelConfig, CreatedAt: g.CreatedAt, UpdatedAt: g.UpdatedAt, } diff --git a/backend/internal/repository/api_key_repo_integration_test.go b/backend/internal/repository/api_key_repo_integration_test.go index 7d5c1826..e926ed86 100644 --- a/backend/internal/repository/api_key_repo_integration_test.go +++ b/backend/internal/repository/api_key_repo_integration_test.go @@ -86,6 +86,45 @@ func (s *APIKeyRepoSuite) TestGetByKey_NotFound() { s.Require().Error(err, "expected error for non-existent key") } +func (s *APIKeyRepoSuite) TestGetByKeyForAuth_PreservesMessagesDispatchModelConfig() { + user := s.mustCreateUser("getbykey-auth-dispatch@test.com") + group, err := s.client.Group.Create(). + SetName("g-auth-dispatch"). + SetPlatform(service.PlatformOpenAI). + SetStatus(service.StatusActive). + SetSubscriptionType(service.SubscriptionTypeStandard). + SetRateMultiplier(1). + SetAllowMessagesDispatch(true). + SetDefaultMappedModel("gpt-5.4"). + SetMessagesDispatchModelConfig(service.OpenAIMessagesDispatchModelConfig{ + OpusMappedModel: "gpt-5.4-nano", + SonnetMappedModel: "gpt-5.3-codex", + HaikuMappedModel: "gpt-5.4-mini", + ExactModelMappings: map[string]string{ + "claude-sonnet-4.5": "gpt-5.4-nano", + }, + }). + Save(s.ctx) + s.Require().NoError(err) + + key := &service.APIKey{ + UserID: user.ID, + Key: "sk-getbykey-auth-dispatch", + Name: "Dispatch Key", + GroupID: &group.ID, + Status: service.StatusActive, + } + s.Require().NoError(s.repo.Create(s.ctx, key)) + + got, err := s.repo.GetByKeyForAuth(s.ctx, key.Key) + s.Require().NoError(err) + s.Require().NotNil(got.Group) + s.Require().True(got.Group.AllowMessagesDispatch) + s.Require().Equal("gpt-5.4", got.Group.DefaultMappedModel) + s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.OpusMappedModel) + s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.ExactModelMappings["claude-sonnet-4.5"]) +} + // --- Update --- func (s *APIKeyRepoSuite) TestUpdate() { diff --git a/backend/internal/repository/api_key_repo_messages_dispatch_unit_test.go b/backend/internal/repository/api_key_repo_messages_dispatch_unit_test.go new file mode 100644 index 00000000..aba62ead --- /dev/null +++ b/backend/internal/repository/api_key_repo_messages_dispatch_unit_test.go @@ -0,0 +1,74 @@ +package repository + +import ( + "context" + "testing" + + dbent "github.com/Wei-Shaw/sub2api/ent" + "github.com/Wei-Shaw/sub2api/internal/service" + "github.com/stretchr/testify/require" +) + +func TestGroupEntityToService_PreservesMessagesDispatchModelConfig(t *testing.T) { + group := &dbent.Group{ + ID: 1, + Name: "openai-dispatch", + Platform: service.PlatformOpenAI, + Status: service.StatusActive, + SubscriptionType: service.SubscriptionTypeStandard, + RateMultiplier: 1, + AllowMessagesDispatch: true, + DefaultMappedModel: "gpt-5.4", + MessagesDispatchModelConfig: service.OpenAIMessagesDispatchModelConfig{ + OpusMappedModel: "gpt-5.4-nano", + SonnetMappedModel: "gpt-5.3-codex", + HaikuMappedModel: "gpt-5.4-mini", + ExactModelMappings: map[string]string{ + "claude-sonnet-4.5": "gpt-5.4-nano", + }, + }, + } + + got := groupEntityToService(group) + require.NotNil(t, got) + require.Equal(t, group.MessagesDispatchModelConfig, got.MessagesDispatchModelConfig) +} + +func TestAPIKeyRepository_GetByKeyForAuth_PreservesMessagesDispatchModelConfig_SQLite(t *testing.T) { + repo, client := newAPIKeyRepoSQLite(t) + ctx := context.Background() + user := mustCreateAPIKeyRepoUser(t, ctx, client, "getbykey-auth-dispatch-unit@test.com") + + group, err := client.Group.Create(). + SetName("g-auth-dispatch-unit"). + SetPlatform(service.PlatformOpenAI). + SetStatus(service.StatusActive). + SetSubscriptionType(service.SubscriptionTypeStandard). + SetRateMultiplier(1). + SetAllowMessagesDispatch(true). + SetDefaultMappedModel("gpt-5.4"). + SetMessagesDispatchModelConfig(service.OpenAIMessagesDispatchModelConfig{ + OpusMappedModel: "gpt-5.4-nano", + SonnetMappedModel: "gpt-5.3-codex", + HaikuMappedModel: "gpt-5.4-mini", + ExactModelMappings: map[string]string{ + "claude-sonnet-4.5": "gpt-5.4-nano", + }, + }). + Save(ctx) + require.NoError(t, err) + + key := &service.APIKey{ + UserID: user.ID, + Key: "sk-getbykey-auth-dispatch-unit", + Name: "Dispatch Key Unit", + GroupID: &group.ID, + Status: service.StatusActive, + } + require.NoError(t, repo.Create(ctx, key)) + + got, err := repo.GetByKeyForAuth(ctx, key.Key) + require.NoError(t, err) + require.NotNil(t, got.Group) + require.Equal(t, group.MessagesDispatchModelConfig, got.Group.MessagesDispatchModelConfig) +} diff --git a/backend/internal/repository/group_repo_integration_test.go b/backend/internal/repository/group_repo_integration_test.go index eccf5cea..f91dae43 100644 --- a/backend/internal/repository/group_repo_integration_test.go +++ b/backend/internal/repository/group_repo_integration_test.go @@ -113,6 +113,33 @@ func (s *GroupRepoSuite) TestUpdate() { s.Require().Equal("updated", got.Name) } +func (s *GroupRepoSuite) TestGetByID_PreservesMessagesDispatchModelConfig() { + group := &service.Group{ + Name: "openai-dispatch", + Platform: service.PlatformOpenAI, + RateMultiplier: 1.0, + IsExclusive: false, + Status: service.StatusActive, + SubscriptionType: service.SubscriptionTypeStandard, + AllowMessagesDispatch: true, + DefaultMappedModel: "gpt-5.4", + MessagesDispatchModelConfig: service.OpenAIMessagesDispatchModelConfig{ + OpusMappedModel: "gpt-5.4", + SonnetMappedModel: "gpt-5.3-codex", + HaikuMappedModel: "gpt-5.4-mini", + ExactModelMappings: map[string]string{ + "claude-sonnet-4.5": "gpt-5.4-nano", + }, + }, + } + + s.Require().NoError(s.repo.Create(s.ctx, group)) + + got, err := s.repo.GetByID(s.ctx, group.ID) + s.Require().NoError(err) + s.Require().Equal(group.MessagesDispatchModelConfig, got.MessagesDispatchModelConfig) +} + func (s *GroupRepoSuite) TestDelete() { group := &service.Group{ Name: "to-delete",