Skip to content
This repository was archived by the owner on Sep 30, 2024. It is now read-only.

Commit cec288d

Browse files
bobheadxibobheadxi
authored
fix/enterpriseportal, fix/codygateway: zero-value durations and missing active licenses (#64378)
This change ensure we correctly handle: 1. In Enterprise Portal, where no active license is available, we return ratelimit=0 intervalduration=0, from the source `PLAN` (as this is determined by the lack of a plan) 2. In Cody Gateway, where intervalduration=0, we do not grant access to that feature ## Test plan Unit tests
1 parent 3edfaa8 commit cec288d

File tree

5 files changed

+104
-21
lines changed

5 files changed

+104
-21
lines changed

‎cmd/cody-gateway/internal/actor/productsubscription/productsubscription.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -324,7 +324,7 @@ func newActor(
324324
Source: source,
325325
}
326326

327-
if rl := s.GetChatCompletionsRateLimit(); rl != nil {
327+
if rl := s.GetChatCompletionsRateLimit(); rl != nil && rl.IntervalDuration.AsDuration() > 0 {
328328
a.RateLimits[codygateway.FeatureChatCompletions] = actor.NewRateLimitWithPercentageConcurrency(
329329
int64(rl.Limit),
330330
rl.IntervalDuration.AsDuration(),
@@ -333,7 +333,7 @@ func newActor(
333333
)
334334
}
335335

336-
if rl := s.GetCodeCompletionsRateLimit(); rl != nil {
336+
if rl := s.GetCodeCompletionsRateLimit(); rl != nil && rl.IntervalDuration.AsDuration() > 0 {
337337
a.RateLimits[codygateway.FeatureCodeCompletions] = actor.NewRateLimitWithPercentageConcurrency(
338338
int64(rl.Limit),
339339
rl.IntervalDuration.AsDuration(),
@@ -342,7 +342,7 @@ func newActor(
342342
)
343343
}
344344

345-
if rl := s.GetEmbeddingsRateLimit(); rl != nil {
345+
if rl := s.GetEmbeddingsRateLimit(); rl != nil && rl.IntervalDuration.AsDuration() > 0 {
346346
a.RateLimits[codygateway.FeatureEmbeddings] = actor.NewRateLimitWithPercentageConcurrency(
347347
int64(rl.Limit),
348348
rl.IntervalDuration.AsDuration(),

‎cmd/cody-gateway/internal/actor/productsubscription/productsubscription_test.go‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,31 @@ func TestNewActor(t *testing.T) {
108108
}
109109
},
110110
"lastUpdated": "2024-06-03T20:03:07-07:00"
111+
}`),
112+
},
113+
{
114+
name: "enabled, rate limit has invalid duration",
115+
args: args{
116+
&codyaccessv1.CodyGatewayAccess{
117+
SubscriptionId: "es_1234uuid",
118+
SubscriptionDisplayName: "My Subscription",
119+
Enabled: true,
120+
EmbeddingsRateLimit: &codyaccessv1.CodyGatewayRateLimit{
121+
Limit: 10,
122+
IntervalDuration: &durationpb.Duration{},
123+
},
124+
},
125+
},
126+
wantActor: autogold.Expect(`{
127+
"key": "sekret_token",
128+
"id": "1234uuid",
129+
"name": "My Subscription",
130+
"accessEnabled": true,
131+
"endpointAccess": {
132+
"/v1/attribution": true
133+
},
134+
"rateLimits": {},
135+
"lastUpdated": "2024-06-03T20:03:07-07:00"
111136
}`),
112137
},
113138
{

‎cmd/enterprise-portal/internal/codyaccessservice/BUILD.bazel‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,8 @@ go_test(
5151
deps = [
5252
"//cmd/enterprise-portal/internal/database/codyaccess",
5353
"//cmd/enterprise-portal/internal/samsm2m",
54+
"//internal/license",
55+
"//internal/licensing",
5456
"//lib/enterpriseportal/codyaccess/v1:codyaccess",
5557
"@com_connectrpc_connect//:connect",
5658
"@com_github_derision_test_go_mockgen_v2//testutil/require",

‎cmd/enterprise-portal/internal/codyaccessservice/adapters.go‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -104,21 +104,26 @@ func maybeApplyOverride[T ~int32 | ~int64](limit *T, overrideValue T, overrideVa
104104
// plan and applying known overrides on top. This closely models the existing
105105
// codyGatewayAccessResolver in 'cmd/frontend/internal/dotcom/productsubscription'.
106106
func evaluateCodyGatewayAccessRateLimits(access *codyaccess.CodyGatewayAccessWithSubscriptionDetails) CodyGatewayRateLimits {
107-
// Set defaults for everything based on active licnese plan and user count.
107+
// Set defaults for everything based on active license plan and user count.
108108
// If there isn't one, zero values apply.
109-
activeLicense := pointers.DerefZero(access.ActiveLicenseInfo)
110-
p := licensing.PlanFromTags(activeLicense.Tags)
111-
userCount := pointers.Ptr(int(activeLicense.UserCount))
112-
113-
limits := CodyGatewayRateLimits{
114-
ChatSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
115-
Chat: licensing.NewCodyGatewayChatRateLimit(p, userCount),
116-
117-
CodeSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
118-
Code: licensing.NewCodyGatewayCodeRateLimit(p, userCount),
119-
120-
EmbeddingsSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
121-
Embeddings: licensing.NewCodyGatewayEmbeddingsRateLimit(p, userCount),
109+
limits := CodyGatewayRateLimits{}
110+
if access.ActiveLicenseInfo != nil {
111+
// Infer defaults for the active license
112+
var (
113+
activeLicense = *access.ActiveLicenseInfo
114+
plan = licensing.PlanFromTags(activeLicense.Tags)
115+
userCount = pointers.Ptr(int(activeLicense.UserCount))
116+
)
117+
limits = CodyGatewayRateLimits{
118+
ChatSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
119+
Chat: licensing.NewCodyGatewayChatRateLimit(plan, userCount),
120+
121+
CodeSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
122+
Code: licensing.NewCodyGatewayCodeRateLimit(plan, userCount),
123+
124+
EmbeddingsSource: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
125+
Embeddings: licensing.NewCodyGatewayEmbeddingsRateLimit(plan, userCount),
126+
}
122127
}
123128

124129
// Chat

‎cmd/enterprise-portal/internal/codyaccessservice/adapters_test.go‎

Lines changed: 55 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,12 @@ import (
66

77
"github.com/hexops/autogold/v2"
88
"github.com/stretchr/testify/assert"
9+
"google.golang.org/protobuf/types/known/durationpb"
910

1011
"github.com/sourcegraph/sourcegraph/cmd/enterprise-portal/internal/database/codyaccess"
12+
"github.com/sourcegraph/sourcegraph/internal/license"
13+
"github.com/sourcegraph/sourcegraph/internal/licensing"
14+
codyaccessv1 "github.com/sourcegraph/sourcegraph/lib/enterpriseportal/codyaccess/v1"
1115
)
1216

1317
func TestConvertAccessAttrsToProto(t *testing.T) {
@@ -43,20 +47,67 @@ func TestConvertAccessAttrsToProto(t *testing.T) {
4347
assert.Empty(t, proto.GetAccessTokens())
4448
})
4549

46-
t.Run("enabled returns everything", func(t *testing.T) {
50+
t.Run("enabled WITHOUT active license returns nothing", func(t *testing.T) {
4751
proto := convertAccessAttrsToProto(&codyaccess.CodyGatewayAccessWithSubscriptionDetails{
4852
CodyGatewayAccess: codyaccess.CodyGatewayAccess{
4953
Enabled: true,
54+
// no overrides
55+
},
56+
ActiveLicenseInfo: nil, // no active license
57+
})
58+
assert.True(t, proto.Enabled)
59+
// Returns non-nil rate limits
60+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
61+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
62+
IntervalDuration: &durationpb.Duration{},
63+
}).Equal(t, proto.ChatCompletionsRateLimit)
64+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
65+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
66+
IntervalDuration: &durationpb.Duration{},
67+
}).Equal(t, proto.CodeCompletionsRateLimit)
68+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
69+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
70+
IntervalDuration: &durationpb.Duration{},
71+
}).Equal(t, proto.EmbeddingsRateLimit)
72+
})
73+
74+
t.Run("enabled with active license returns plan defaults", func(t *testing.T) {
75+
proto := convertAccessAttrsToProto(&codyaccess.CodyGatewayAccessWithSubscriptionDetails{
76+
CodyGatewayAccess: codyaccess.CodyGatewayAccess{
77+
Enabled: true,
78+
// no overrides
79+
},
80+
ActiveLicenseInfo: &license.Info{
81+
UserCount: 10,
82+
Tags: []string{licensing.PlanCodyEnterprise.Tag()},
5083
},
5184
LicenseKeyHashes: [][]byte{[]byte("abc"), []byte("efg")},
5285
})
5386
assert.True(t, proto.Enabled)
5487
// NOTE: These are not real access tokens
5588
autogold.Expect([]string{`token:"slk_616263"`, `token:"slk_656667"`}).Equal(t, toStrings(proto.GetAccessTokens()))
5689
// Returns non-nil rate limits
57-
assert.NotNil(t, proto.ChatCompletionsRateLimit)
58-
assert.NotNil(t, proto.CodeCompletionsRateLimit)
59-
assert.NotNil(t, proto.EmbeddingsRateLimit)
90+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
91+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
92+
Limit: 100,
93+
IntervalDuration: &durationpb.Duration{
94+
Seconds: 86400,
95+
},
96+
}).Equal(t, proto.ChatCompletionsRateLimit)
97+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
98+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
99+
Limit: 10000,
100+
IntervalDuration: &durationpb.Duration{
101+
Seconds: 86400,
102+
},
103+
}).Equal(t, proto.CodeCompletionsRateLimit)
104+
autogold.Expect(&codyaccessv1.CodyGatewayRateLimit{
105+
Source: codyaccessv1.CodyGatewayRateLimitSource_CODY_GATEWAY_RATE_LIMIT_SOURCE_PLAN,
106+
Limit: 33333333,
107+
IntervalDuration: &durationpb.Duration{
108+
Seconds: 86400,
109+
},
110+
}).Equal(t, proto.EmbeddingsRateLimit)
60111
})
61112
}
62113

0 commit comments

Comments
 (0)