Skip to content

week #29

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sergei-bronnikov wants to merge 4 commits into
base: main
Choose a base branch
from
Open

week #29

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4c146b5
week
sergei-bronnikov Aug 20, 2026
079f1ee
wip
sergei-bronnikov Aug 21, 2026
4920dfa
count. validate
sergei-bronnikov Aug 21, 2026
599680a
check duplicated
sergei-bronnikov Aug 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
196 changes: 126 additions & 70 deletions internal/key/key.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package key
import (
"errors"
"fmt"
"slices"
"strings"
"time"

Expand All @@ -11,27 +12,41 @@ import (

const RevokedReasonExpired string = "expired"

type ExtendedBudgetLimitItem struct {
LimitInUsdOverTime float64 `json:"costLimitOverTime"`
Unit TimeUnit `json:"unit"`
}

func (i *ExtendedBudgetLimitItem) ExtendedKey(k string) string {
return fmt.Sprintf("%s-%s", k, i.Unit)
}

type ExtendedBudgetLimit struct {
Items []ExtendedBudgetLimitItem `json:"items"`
}

type UpdateKey struct {
Name string `json:"name"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
Revoked *bool `json:"revoked"`
RevokedReason string `json:"revokedReason"`
Key string `json:"key"`
SettingId string `json:"settingId"`
SettingIds []string `json:"settingIds"`
CostLimitInUsd *float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime *int `json:"rateLimitOverTime"`
RateLimitUnit *TimeUnit `json:"rateLimitUnit"`
RequestsLimit *int `json:"requestsLimit"`
AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"`
ShouldLogRequest *bool `json:"shouldLogRequest"`
ShouldLogResponse *bool `json:"shouldLogResponse"`
RotationEnabled *bool `json:"rotationEnabled"`
PolicyId *string `json:"policyId"`
IsKeyNotHashed *bool `json:"isKeyNotHashed"`
Name string `json:"name"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
Revoked *bool `json:"revoked"`
RevokedReason string `json:"revokedReason"`
Key string `json:"key"`
SettingId string `json:"settingId"`
SettingIds []string `json:"settingIds"`
CostLimitInUsd *float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime *float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit *TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime *int `json:"rateLimitOverTime"`
RateLimitUnit *TimeUnit `json:"rateLimitUnit"`
RequestsLimit *int `json:"requestsLimit"`
AllowedPaths *[]PathConfig `json:"allowedPaths,omitempty"`
ShouldLogRequest *bool `json:"shouldLogRequest"`
ShouldLogResponse *bool `json:"shouldLogResponse"`
RotationEnabled *bool `json:"rotationEnabled"`
PolicyId *string `json:"policyId"`
IsKeyNotHashed *bool `json:"isKeyNotHashed"`
ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"`
}

func (uk *UpdateKey) Validate() error {
Expand Down Expand Up @@ -125,7 +140,7 @@ func (uk *UpdateKey) Validate() error {
return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified")
}

if *uk.RateLimitOverTime != 0 && *uk.RateLimitUnit != HourTimeUnit && *uk.RateLimitUnit != MinuteTimeUnit && *uk.RateLimitUnit != SecondTimeUnit && *uk.RateLimitUnit != DayTimeUnit {
if *uk.RateLimitOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.RateLimitUnit) {
return internal_errors.NewValidationError("rate limit unit can not be identified")
}
}
Expand All @@ -143,7 +158,7 @@ func (uk *UpdateKey) Validate() error {
return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified")
}

if *uk.CostLimitInUsdOverTime != 0 && *uk.CostLimitInUsdUnit != DayTimeUnit && *uk.CostLimitInUsdUnit != HourTimeUnit && *uk.CostLimitInUsdUnit != MonthTimeUnit && *uk.CostLimitInUsdUnit != MinuteTimeUnit {
if *uk.CostLimitInUsdOverTime != 0 && !slices.Contains(AllowedTimeUnits, *uk.CostLimitInUsdUnit) {
return internal_errors.NewValidationError("cost limit unit can not be identified")
}
}
Expand All @@ -157,28 +172,29 @@ type PathConfig struct {
}

type RequestKey struct {
Name string `json:"name"`
CreatedAt int64 `json:"createdAt"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
KeyId string `json:"keyId"`
Key string `json:"key"`
CostLimitInUsd float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime int `json:"rateLimitOverTime"`
RateLimitUnit TimeUnit `json:"rateLimitUnit"`
Ttl string `json:"ttl"`
KeyRing string `json:"keyRing"`
SettingId string `json:"settingId"`
AllowedPaths []PathConfig `json:"allowedPaths"`
SettingIds []string `json:"settingIds"`
ShouldLogRequest bool `json:"shouldLogRequest"`
ShouldLogResponse bool `json:"shouldLogResponse"`
RotationEnabled bool `json:"rotationEnabled"`
PolicyId string `json:"policyId"`
IsKeyNotHashed bool `json:"isKeyNotHashed"`
RequestsLimit int `json:"requestsLimit"`
Name string `json:"name"`
CreatedAt int64 `json:"createdAt"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
KeyId string `json:"keyId"`
Key string `json:"key"`
CostLimitInUsd float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime int `json:"rateLimitOverTime"`
RateLimitUnit TimeUnit `json:"rateLimitUnit"`
Ttl string `json:"ttl"`
KeyRing string `json:"keyRing"`
SettingId string `json:"settingId"`
AllowedPaths []PathConfig `json:"allowedPaths"`
SettingIds []string `json:"settingIds"`
ShouldLogRequest bool `json:"shouldLogRequest"`
ShouldLogResponse bool `json:"shouldLogResponse"`
RotationEnabled bool `json:"rotationEnabled"`
PolicyId string `json:"policyId"`
IsKeyNotHashed bool `json:"isKeyNotHashed"`
RequestsLimit int `json:"requestsLimit"`
ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"`
}

func (rk *RequestKey) Validate() error {
Expand Down Expand Up @@ -285,7 +301,7 @@ func (rk *RequestKey) Validate() error {
return internal_errors.NewValidationError("rate limit unit can not be empty if rate limit over time is specified")
}

if rk.RateLimitUnit != HourTimeUnit && rk.RateLimitUnit != MinuteTimeUnit && rk.RateLimitUnit != SecondTimeUnit && rk.RateLimitUnit != DayTimeUnit {
if !slices.Contains(AllowedTimeUnits, rk.RateLimitUnit) {
return internal_errors.NewValidationError("rate limit unit can not be identified")
}
}
Expand All @@ -295,11 +311,40 @@ func (rk *RequestKey) Validate() error {
return internal_errors.NewValidationError("cost limit unit can not be empty if cost limit over time is specified")
}

if rk.CostLimitInUsdUnit != DayTimeUnit && rk.CostLimitInUsdUnit != HourTimeUnit && rk.CostLimitInUsdUnit != MonthTimeUnit && rk.CostLimitInUsdUnit != MinuteTimeUnit {
if !slices.Contains(AllowedTimeUnits, rk.CostLimitInUsdUnit) {
return internal_errors.NewValidationError("cost limit unit can not be identified")
}
}
if err := validateExtendedBudgetLimit(rk.ExtendedBudgetLimit); err != nil {
return err
}
return nil
}

func validateExtendedBudgetLimit(ebl *ExtendedBudgetLimit) error {
if ebl == nil {
return nil
}

seenUnits := make(map[TimeUnit]struct{}, len(ebl.Items))
for index, item := range ebl.Items {
if item.LimitInUsdOverTime < 0 {
return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].costLimitOverTime is invalid", index))
}

if len(item.Unit) == 0 {
return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is invalid", index))
}

if !slices.Contains(AllowedTimeUnits, item.Unit) {
return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit can not be identified", index))
}

if _, exists := seenUnits[item.Unit]; exists {
return internal_errors.NewValidationError(fmt.Sprintf("extendedBudgetLimit.items[%d].unit is duplicated", index))
}
seenUnits[item.Unit] = struct{}{}
}
return nil
}

Expand All @@ -310,34 +355,45 @@ const (
MinuteTimeUnit TimeUnit = "m"
SecondTimeUnit TimeUnit = "s"
DayTimeUnit TimeUnit = "d"
WeekTimeUnit TimeUnit = "w"
MonthTimeUnit TimeUnit = "mo"
)

var AllowedTimeUnits = []TimeUnit{
HourTimeUnit,
MinuteTimeUnit,
SecondTimeUnit,
DayTimeUnit,
WeekTimeUnit,
MonthTimeUnit,
}

type ResponseKey struct {
Name string `json:"name"`
CreatedAt int64 `json:"createdAt"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
KeyId string `json:"keyId"`
Revoked bool `json:"revoked"`
Key string `json:"key"`
RevokedReason string `json:"revokedReason"`
CostLimitInUsd float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime int `json:"rateLimitOverTime"`
RateLimitUnit TimeUnit `json:"rateLimitUnit"`
RequestsLimit int `json:"requestsLimit"`
Ttl string `json:"ttl"`
KeyRing string `json:"keyRing"`
SettingId string `json:"settingId"`
AllowedPaths []PathConfig `json:"allowedPaths"`
SettingIds []string `json:"settingIds"`
ShouldLogRequest bool `json:"shouldLogRequest"`
ShouldLogResponse bool `json:"shouldLogResponse"`
RotationEnabled bool `json:"rotationEnabled"`
PolicyId string `json:"policyId"`
IsKeyNotHashed bool `json:"isKeyNotHashed"`
Name string `json:"name"`
CreatedAt int64 `json:"createdAt"`
UpdatedAt int64 `json:"updatedAt"`
Tags []string `json:"tags"`
KeyId string `json:"keyId"`
Revoked bool `json:"revoked"`
Key string `json:"key"`
RevokedReason string `json:"revokedReason"`
CostLimitInUsd float64 `json:"costLimitInUsd"`
CostLimitInUsdOverTime float64 `json:"costLimitInUsdOverTime"`
CostLimitInUsdUnit TimeUnit `json:"costLimitInUsdUnit"`
RateLimitOverTime int `json:"rateLimitOverTime"`
RateLimitUnit TimeUnit `json:"rateLimitUnit"`
RequestsLimit int `json:"requestsLimit"`
Ttl string `json:"ttl"`
KeyRing string `json:"keyRing"`
SettingId string `json:"settingId"`
AllowedPaths []PathConfig `json:"allowedPaths"`
SettingIds []string `json:"settingIds"`
ShouldLogRequest bool `json:"shouldLogRequest"`
ShouldLogResponse bool `json:"shouldLogResponse"`
RotationEnabled bool `json:"rotationEnabled"`
PolicyId string `json:"policyId"`
IsKeyNotHashed bool `json:"isKeyNotHashed"`
ExtendedBudgetLimit *ExtendedBudgetLimit `json:"extendedBudgetLimit"`
}

func (rk *ResponseKey) GetSettingIds() []string {
Expand Down
57 changes: 57 additions & 0 deletions internal/key/key_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
package key

import "testing"

func TestValidateExtendedBudgetLimit(t *testing.T) {
tests := []struct {
name string
ebl *ExtendedBudgetLimit
wantErr string
}{
{
name: "nil limit is valid",
ebl: nil,
wantErr: "",
},
{
name: "duplicate units are rejected",
ebl: &ExtendedBudgetLimit{
Items: []ExtendedBudgetLimitItem{
{LimitInUsdOverTime: 1, Unit: HourTimeUnit},
{LimitInUsdOverTime: 2, Unit: HourTimeUnit},
},
},
wantErr: "extendedBudgetLimit.items[1].unit is duplicated",
},
{
name: "unique units are valid",
ebl: &ExtendedBudgetLimit{
Items: []ExtendedBudgetLimitItem{
{LimitInUsdOverTime: 1, Unit: HourTimeUnit},
{LimitInUsdOverTime: 2, Unit: MinuteTimeUnit},
},
},
wantErr: "",
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
err := validateExtendedBudgetLimit(tt.ebl)
if tt.wantErr == "" {
if err != nil {
t.Fatalf("expected no error, got %v", err)
}
return
}

if err == nil {
t.Fatalf("expected error %q, got nil", tt.wantErr)
}

if err.Error() != tt.wantErr {
t.Fatalf("expected error %q, got %q", tt.wantErr, err.Error())
}
})
}
}
2 changes: 1 addition & 1 deletion internal/message/consumer.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ type Consumer struct {
}

type recorder interface {
RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error
RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error
RecordUserSpend(userId string, micros int64, costLimitUnit key.TimeUnit) error
RecordEvent(e *event.Event) error
RecordKeyRequestSpent(keyId string) error
Expand Down
2 changes: 1 addition & 1 deletion internal/message/handler.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -340,7 +340,7 @@ func (h *Handler) HandleEventWithRequestAndResponse(m Message) error {

if e.Event.CostInUsd != 0 {
micros := int64(e.Event.CostInUsd * 1000000)
err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit)
err = h.recorder.RecordKeySpend(e.Event.KeyId, micros, e.Key.CostLimitInUsdUnit, e.Key.ExtendedBudgetLimit)
if err != nil {
telemetry.Incr("bricksllm.message.handler.handle_event_with_request_and_response.record_key_spend_error", nil, 1)
h.log.Debug("error when recording key spend", zap.Error(err))
Expand Down
11 changes: 10 additions & 1 deletion internal/recorder/recorder.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,7 @@ func (r *Recorder) RecordUserSpend(userId string, micros int64, costLimitUnit ke
return nil
}

func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit) error {
func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.TimeUnit, extendedLimits *key.ExtendedBudgetLimit) error {
err := r.s.IncrementCounter(keyId, micros)
if err != nil {
return err
Expand All @@ -73,6 +73,15 @@ func (r *Recorder) RecordKeySpend(keyId string, micros int64, costLimitUnit key.
}
}

if extendedLimits != nil {
for _, item := range extendedLimits.Items {
err = r.c.IncrementCounter(item.ExtendedKey(keyId), item.Unit, int64(micros))
if err != nil {
return err
}
}
}

return nil
}

Expand Down
Loading
Loading