From 77578169432958a04f696154a73027726f2f2ff0 Mon Sep 17 00:00:00 2001 From: Jawad Date: Thu, 24 Sep 2026 10:43:53 -0500 Subject: [PATCH] feat(paymodels): add auto-select-single-paymodel to survive workspace terminate Terminating a workspace runs resetCurrentPaymodel in a goroutine, which sets current_pay_model=false on *every* row for that user. On commons where users are never asked to pick a pay model, that leaves them unable to launch again: getCurrentPayModel finds no row with current=true, sees that active rows do exist, and returns nil, so /launch responds 500 with "Current Paymodel is not set. Launch forbidden". The default-pay-model fallback immediately below is unreachable in this case -- it only fires for users with zero rows in the table. Add an opt-in "auto-select-single-paymodel" config key. When set, a user whose only active pay model is not flagged current has that row returned as their current one. The row is returned as-is rather than substituting DefaultPayModel so it keeps its bmh_workspace_id: pods are annotated with the right paymodel_type and updatePayModelCost can still find the row to bill, and the row's own hard/soft limits and accrued total-usage apply. Users with more than one pay model are untouched -- they have a real choice to make and the flag must not make it for them. Default is off, so behaviour is unchanged for every commons that does not opt in. --- hatchery/config.go | 42 +++++++++++++---------- hatchery/paymodels.go | 7 ++++ hatchery/paymodels_test.go | 68 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 99 insertions(+), 18 deletions(-) diff --git a/hatchery/config.go b/hatchery/config.go index d2e5e50d..a6c4a586 100644 --- a/hatchery/config.go +++ b/hatchery/config.go @@ -178,24 +178,30 @@ type HatcheryConfig struct { DefaultPayModel PayModel `json:"default-pay-model"` HashedUsernames bool `json:"hashed-usernames"` // DisableLocalWS bool `json:"disable-local-ws"` - SkipNodeSelector bool `json:"skip-node-selector"` - UseInteralServicesURL bool `json:"use-internal-services-url"` - PayModels []PayModel `json:"pay-models"` - PayModelsDynamodbTable string `json:"pay-models-dynamodb-table"` - PayModelsDynamodbArn string `json:"pay-models-dynamodb-arn"` - LicenseUserMapsTable string `json:"license-user-maps-dynamodb-table"` - LicenseUserMapsGSI string `json:"license-user-maps-global-secondary-index"` - License LicenseInfo `json:"license"` - SubDir string `json:"sub-dir"` - Containers []Container `json:"containers"` - UserVolumeSize string `json:"user-volume-size"` - Sidecar SidecarContainer `json:"sidecar"` - MoreConfigs []AppConfigInfo `json:"more-configs"` - PrismaConfig PrismaConfig `json:"prisma"` - S3Config S3Config `json:"s3-config"` - NextflowGlobalConfig NextflowGlobalConfig `json:"nextflow-global"` - Pricing Pricing `json:"pricing"` - SharedWorkspace SharedWorkspaceConfig `json:"shared-workspace"` + SkipNodeSelector bool `json:"skip-node-selector"` + UseInteralServicesURL bool `json:"use-internal-services-url"` + PayModels []PayModel `json:"pay-models"` + PayModelsDynamodbTable string `json:"pay-models-dynamodb-table"` + PayModelsDynamodbArn string `json:"pay-models-dynamodb-arn"` + // AutoSelectSinglePayModel treats a user's only active pay model as their + // current one, even when current_pay_model is false in DynamoDB. Terminating + // a workspace clears that flag for every row (see resetCurrentPaymodel), so on + // commons where users never pick a pay model explicitly this otherwise leaves + // them unable to launch until someone flips the flag by hand. + AutoSelectSinglePayModel bool `json:"auto-select-single-paymodel"` + LicenseUserMapsTable string `json:"license-user-maps-dynamodb-table"` + LicenseUserMapsGSI string `json:"license-user-maps-global-secondary-index"` + License LicenseInfo `json:"license"` + SubDir string `json:"sub-dir"` + Containers []Container `json:"containers"` + UserVolumeSize string `json:"user-volume-size"` + Sidecar SidecarContainer `json:"sidecar"` + MoreConfigs []AppConfigInfo `json:"more-configs"` + PrismaConfig PrismaConfig `json:"prisma"` + S3Config S3Config `json:"s3-config"` + NextflowGlobalConfig NextflowGlobalConfig `json:"nextflow-global"` + Pricing Pricing `json:"pricing"` + SharedWorkspace SharedWorkspaceConfig `json:"shared-workspace"` // OIDCProviderARN is the full ARN of the EKS cluster's OIDC provider, used to // build IRSA trust policies. It is a cluster-level property shared by the // shared-workspace and squashfs software-library features. diff --git a/hatchery/paymodels.go b/hatchery/paymodels.go index b34ddee9..e2119381 100644 --- a/hatchery/paymodels.go +++ b/hatchery/paymodels.go @@ -109,6 +109,13 @@ var getCurrentPayModel = func(userName string) (result *PayModel, err error) { activePayModels, _ := payModelsFromDatabase(userName, false) if activePayModels != nil && len(*activePayModels) > 0 { + // Terminating a workspace resets current_pay_model to false on every + // row, so a user with exactly one pay model is left with nothing + // current and cannot launch until they re-select it. Where that + // choice is not surfaced to users, treat the only option as current. + if Config.Config.AutoSelectSinglePayModel && len(*activePayModels) == 1 { + return &(*activePayModels)[0], nil + } // return nil since there is no current paymodel set by the user return nil, nil } diff --git a/hatchery/paymodels_test.go b/hatchery/paymodels_test.go index 634cfdcc..f8adf214 100644 --- a/hatchery/paymodels_test.go +++ b/hatchery/paymodels_test.go @@ -56,6 +56,13 @@ func Test_GetCurrentPayModel(t *testing.T) { }, } + configWithAutoSelectSinglePayModel := &FullHatcheryConfig{ + Config: HatcheryConfig{ + PayModelsDynamodbTable: "random_non_empty_string", + AutoSelectSinglePayModel: true, + }, + } + defaultPayModelForTest := &PayModel{ Name: "Trial Workspace", Local: true, @@ -146,6 +153,67 @@ func Test_GetCurrentPayModel(t *testing.T) { mockPayModelsFromDB: []PayModel{}, mockDefaultPaymodel: defaultPayModelForTest, }, + { + // A terminate resets current_pay_model on the user's only row. With + // the flag on, that row is still returned so the user can launch. + name: "AutoSelectSingle_SingleActiveNotCurrent", + want: &PayModel{ + Id: "#1", + Name: "Trial Workspace", + CurrentPayModel: false, + Status: "active", + }, + mockConfig: configWithAutoSelectSinglePayModel, + mockCurrentPayModelFromDB: []PayModel{}, + mockPayModelsFromDB: []PayModel{ + { + Id: "#1", + Name: "Trial Workspace", + CurrentPayModel: false, + Status: "active", + }, + }, + mockDefaultPaymodel: defaultPayModelForTest, + }, + { + // With more than one option the user genuinely has a choice to make, + // so the flag must not pick one for them. + name: "AutoSelectSingle_MultipleActiveNotCurrent", + want: nil, + mockConfig: configWithAutoSelectSinglePayModel, + mockCurrentPayModelFromDB: []PayModel{}, + mockPayModelsFromDB: []PayModel{ + { + Id: "#1", + Name: "Trial Workspace", + CurrentPayModel: false, + Status: "active", + }, + { + Id: "#2", + Name: "Direct Pay", + CurrentPayModel: false, + Status: "active", + }, + }, + mockDefaultPaymodel: defaultPayModelForTest, + }, + { + // Same single-row shape as above but flag off: unchanged behaviour. + name: "AutoSelectSingleDisabled_SingleActiveNotCurrent", + want: nil, + mockConfig: configWithDbTable, + mockCurrentPayModelFromDB: []PayModel{}, + mockPayModelsFromDB: []PayModel{ + { + Id: "#1", + Name: "Trial Workspace", + CurrentPayModel: false, + Status: "active", + }, + }, + mockDefaultPaymodel: defaultPayModelForTest, + }, } for _, testcase := range testCases {