Merge pull request #50 from simonsmh/feature/kiro-rate-limiter-config

feat(kiro): rate limiter config, system prompt control, Claude format fixes

Author: kaitranntt

cmd/server/main.go

@@ -557,32 +557,44 @@ func main() {
 		// and don't share config with StartService (which is in the else branch)
 		setKiroIncognitoMode(cfg, useIncognito, noIncognito)
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroLogin(cfg, options)
 	} else if kiroGoogleLogin {
 		// For Kiro auth, default to incognito mode for multi-account support
 		// Users can explicitly override with --no-incognito
 		// Note: This config mutation is safe - auth commands exit after completion
 		setKiroIncognitoMode(cfg, useIncognito, noIncognito)
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroGoogleLogin(cfg, options)
 	} else if kiroAWSLogin {
 		// For Kiro auth, default to incognito mode for multi-account support
 		// Users can explicitly override with --no-incognito
 		setKiroIncognitoMode(cfg, useIncognito, noIncognito)
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroAWSLogin(cfg, options)
 	} else if kiroAWSAuthCode {
 		// For Kiro auth with authorization code flow (better UX)
 		setKiroIncognitoMode(cfg, useIncognito, noIncognito)
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroAWSAuthCodeLogin(cfg, options)
 	} else if kiroImport {
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroImport(cfg, options)
 	} else if kiroIDCLogin {
 		// For Kiro IDC auth, default to incognito mode for multi-account support
 		setKiroIncognitoMode(cfg, useIncognito, noIncognito)
 		kiro.InitFingerprintConfig(cfg)
+		kiro.InitRateLimiterConfig(cfg)
+		kiro.InitSystemPromptInjectConfig(cfg)
 		cmd.DoKiroIDCLogin(cfg, options, kiroIDCStartURL, kiroIDCRegion, kiroIDCFlow)
 	} else {
 		// In cloud deploy mode without config file, just wait for shutdown signals
@@ -680,6 +692,8 @@ func main() {
 			}
 
 			if cfg.AuthDir != "" {
+				kiro.InitRateLimiterConfig(cfg)
+				kiro.InitSystemPromptInjectConfig(cfg)
 				kiro.InitializeAndStart(cfg.AuthDir, cfg)
 				defer kiro.StopGlobalRefreshManager()
 			}

config.example.yaml

@@ -249,6 +249,27 @@ nonstream-keepalive-interval: 0
 #    profile-arn: "arn:aws:codewhisperer:us-east-1:..."
 #    proxy-url: "socks5://proxy.example.com:1080" # optional: proxy override
 
+# Kiro rate limiter configuration
+# Controls request throttling and backoff behavior for Kiro tokens.
+# All fields are optional; omitted fields fall back to defaults shown below.
+#kiro-rate-limit:
+#  enabled: true                # set to true to enable rate limiting (default: false)
+#  min-token-interval: "1s"     # minimum interval between requests per token (default: 1s)
+#  max-token-interval: "2s"     # maximum interval between requests per token (default: 2s)
+#  daily-max-requests: 500      # maximum requests per token per day (default: 500)
+#  jitter-percent: 0.3          # random jitter ratio applied to intervals (default: 0.3)
+#  backoff-base: "30s"          # base duration for exponential backoff on failure (default: 30s)
+#  backoff-max: "5m"            # maximum backoff duration cap (default: 5m)
+#  backoff-multiplier: 1.5      # exponential backoff multiplier (default: 1.5)
+#  suspend-cooldown: "1h"       # cooldown after suspension keywords detected (default: 1h)
+
+# Controls whether system prompts are wrapped with --- SYSTEM PROMPT --- markers
+# and injected into Kiro user messages. Kiro/Amazon Q API has no native system
+# field, so when enabled the proxy prepends the wrapped system prompt to the
+# user message. By default the proxy drops system prompts entirely.
+# Default: false.
+#kiro-system-prompt-inject-enable: true
+
 # Kilocode (OAuth-based code assistant)
 # Note: Kilocode uses OAuth device flow authentication.
 # Use the CLI command: ./server --kilo-login

internal/auth/kiro/rate_limiter.go

@@ -6,6 +6,8 @@ import (
 	"strings"
 	"sync"
 	"time"
+
+	log "github.com/sirupsen/logrus"
 )
 
 const (
@@ -19,7 +21,7 @@ const (
 	DefaultSuspendCooldown   = 1 * time.Hour
 )
 
-// TokenState Token 状态
+// TokenState tracks per-token request and cooldown state.
 type TokenState struct {
 	LastRequest    time.Time
 	RequestCount   int
@@ -32,10 +34,11 @@ type TokenState struct {
 	SuspendReason  string
 }
 
-// RateLimiter 频率限制器
+// RateLimiter throttles Kiro requests and tracks token health.
 type RateLimiter struct {
 	mu                sync.RWMutex
 	states            map[string]*TokenState
+	enabled           bool
 	minTokenInterval  time.Duration
 	maxTokenInterval  time.Duration
 	dailyMaxRequests  int
@@ -47,24 +50,19 @@ type RateLimiter struct {
 	rng               *rand.Rand
 }
 
-// NewRateLimiter 创建默认配置的频率限制器
+// NewRateLimiter creates a rate limiter with default settings.
 func NewRateLimiter() *RateLimiter {
-	return &RateLimiter{
-		states:            make(map[string]*TokenState),
-		minTokenInterval:  DefaultMinTokenInterval,
-		maxTokenInterval:  DefaultMaxTokenInterval,
-		dailyMaxRequests:  DefaultDailyMaxRequests,
-		jitterPercent:     DefaultJitterPercent,
-		backoffBase:       DefaultBackoffBase,
-		backoffMax:        DefaultBackoffMax,
-		backoffMultiplier: DefaultBackoffMultiplier,
-		suspendCooldown:   DefaultSuspendCooldown,
-		rng:               rand.New(rand.NewSource(time.Now().UnixNano())),
+	rl := &RateLimiter{
+		states: make(map[string]*TokenState),
+		rng:    rand.New(rand.NewSource(time.Now().UnixNano())),
 	}
+	rl.resetConfigLocked()
+	return rl
 }
 
-// RateLimiterConfig 频率限制器配置
+// RateLimiterConfig defines rate limiter settings.
 type RateLimiterConfig struct {
+	Enabled           *bool
 	MinTokenInterval  time.Duration
 	MaxTokenInterval  time.Duration
 	DailyMaxRequests  int
@@ -75,9 +73,37 @@ type RateLimiterConfig struct {
 	SuspendCooldown   time.Duration
 }
 
-// NewRateLimiterWithConfig 使用自定义配置创建频率限制器
+// NewRateLimiterWithConfig creates a rate limiter with custom settings.
 func NewRateLimiterWithConfig(cfg RateLimiterConfig) *RateLimiter {
 	rl := NewRateLimiter()
+	rl.applyConfigLocked(cfg)
+	return rl
+}
+
+// ApplyConfig reapplies config defaults and then applies overrides.
+func (rl *RateLimiter) ApplyConfig(cfg RateLimiterConfig) {
+	rl.mu.Lock()
+	defer rl.mu.Unlock()
+	rl.applyConfigLocked(cfg)
+}
+
+func (rl *RateLimiter) resetConfigLocked() {
+	rl.enabled = false
+	rl.minTokenInterval = DefaultMinTokenInterval
+	rl.maxTokenInterval = DefaultMaxTokenInterval
+	rl.dailyMaxRequests = DefaultDailyMaxRequests
+	rl.jitterPercent = DefaultJitterPercent
+	rl.backoffBase = DefaultBackoffBase
+	rl.backoffMax = DefaultBackoffMax
+	rl.backoffMultiplier = DefaultBackoffMultiplier
+	rl.suspendCooldown = DefaultSuspendCooldown
+}
+
+func (rl *RateLimiter) applyConfigLocked(cfg RateLimiterConfig) {
+	rl.resetConfigLocked()
+	if cfg.Enabled != nil {
+		rl.enabled = *cfg.Enabled
+	}
 	if cfg.MinTokenInterval > 0 {
 		rl.minTokenInterval = cfg.MinTokenInterval
 	}
@@ -102,10 +128,16 @@ func NewRateLimiterWithConfig(cfg RateLimiterConfig) *RateLimiter {
 	if cfg.SuspendCooldown > 0 {
 		rl.suspendCooldown = cfg.SuspendCooldown
 	}
-	return rl
+
+	// Validate interval bounds: max must be > min to avoid rand.Int63n panic.
+	if rl.maxTokenInterval <= rl.minTokenInterval {
+		log.Warnf("kiro: rate limiter max-token-interval (%v) <= min-token-interval (%v), clamping max to min+1s",
+			rl.maxTokenInterval, rl.minTokenInterval)
+		rl.maxTokenInterval = rl.minTokenInterval + time.Second
+	}
 }
 
-// getOrCreateState 获取或创建 Token 状态
+// getOrCreateState returns the existing token state or creates one.
 func (rl *RateLimiter) getOrCreateState(tokenKey string) *TokenState {
 	state, exists := rl.states[tokenKey]
 	if !exists {
@@ -117,7 +149,7 @@ func (rl *RateLimiter) getOrCreateState(tokenKey string) *TokenState {
 	return state
 }
 
-// resetDailyIfNeeded 如果需要则重置每日计数
+// resetDailyIfNeeded resets the daily counter when the window rolls over.
 func (rl *RateLimiter) resetDailyIfNeeded(state *TokenState) {
 	now := time.Now()
 	if now.After(state.DailyResetTime) {
@@ -126,22 +158,26 @@ func (rl *RateLimiter) resetDailyIfNeeded(state *TokenState) {
 	}
 }
 
-// calculateInterval 计算带抖动的随机间隔
+// calculateInterval returns a randomized interval with jitter applied.
 func (rl *RateLimiter) calculateInterval() time.Duration {
 	baseInterval := rl.minTokenInterval + time.Duration(rl.rng.Int63n(int64(rl.maxTokenInterval-rl.minTokenInterval)))
 	jitter := time.Duration(float64(baseInterval) * rl.jitterPercent * (rl.rng.Float64()*2 - 1))
 	return baseInterval + jitter
 }
 
-// WaitForToken 等待 Token 可用（带抖动的随机间隔）
+// WaitForToken blocks until the token is allowed to send another request.
 func (rl *RateLimiter) WaitForToken(tokenKey string) {
+	if !rl.enabled {
+		return
+	}
+
 	rl.mu.Lock()
 	state := rl.getOrCreateState(tokenKey)
 	rl.resetDailyIfNeeded(state)
 
 	now := time.Now()
 
-	// 检查是否在冷却期
+	// Wait until the cooldown window ends.
 	if now.Before(state.CooldownEnd) {
 		waitTime := state.CooldownEnd.Sub(now)
 		rl.mu.Unlock()
@@ -151,7 +187,7 @@ func (rl *RateLimiter) WaitForToken(tokenKey string) {
 		now = time.Now()
 	}
 
-	// 计算距离上次请求的间隔
+	// Respect the randomized spacing since the last request.
 	interval := rl.calculateInterval()
 	nextAllowedTime := state.LastRequest.Add(interval)
 
@@ -169,8 +205,12 @@ func (rl *RateLimiter) WaitForToken(tokenKey string) {
 	rl.mu.Unlock()
 }
 
-// MarkTokenFailed 标记 Token 失败
+// MarkTokenFailed records a failed request for a token.
 func (rl *RateLimiter) MarkTokenFailed(tokenKey string) {
+	if !rl.enabled {
+		return
+	}
+
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
@@ -179,8 +219,12 @@ func (rl *RateLimiter) MarkTokenFailed(tokenKey string) {
 	state.CooldownEnd = time.Now().Add(rl.calculateBackoff(state.FailCount))
 }
 
-// MarkTokenSuccess 标记 Token 成功
+// MarkTokenSuccess clears any failure state after a successful request.
 func (rl *RateLimiter) MarkTokenSuccess(tokenKey string) {
+	if !rl.enabled {
+		return
+	}
+
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
@@ -189,8 +233,12 @@ func (rl *RateLimiter) MarkTokenSuccess(tokenKey string) {
 	state.CooldownEnd = time.Time{}
 }
 
-// CheckAndMarkSuspended 检测暂停错误并标记
+// CheckAndMarkSuspended detects suspension-like errors and records them.
 func (rl *RateLimiter) CheckAndMarkSuspended(tokenKey string, errorMsg string) bool {
+	if !rl.enabled {
+		return false
+	}
+
 	suspendKeywords := []string{
 		"suspended",
 		"banned",
@@ -219,8 +267,12 @@ func (rl *RateLimiter) CheckAndMarkSuspended(tokenKey string, errorMsg string) b
 	return false
 }
 
-// IsTokenAvailable 检查 Token 是否可用
+// IsTokenAvailable reports whether the token can currently be used.
 func (rl *RateLimiter) IsTokenAvailable(tokenKey string) bool {
+	if !rl.enabled {
+		return true
+	}
+
 	rl.mu.RLock()
 	defer rl.mu.RUnlock()
 
@@ -231,20 +283,20 @@ func (rl *RateLimiter) IsTokenAvailable(tokenKey string) bool {
 
 	now := time.Now()
 
-	// 检查是否被暂停
+	// Suspended tokens stay unavailable until their suspension window expires.
 	if state.IsSuspended {
 		if now.After(state.SuspendedAt.Add(rl.suspendCooldown)) {
 			return true
 		}
 		return false
 	}
 
-	// 检查是否在冷却期
+	// Cooldown blocks reuse after failures.
 	if now.Before(state.CooldownEnd) {
 		return false
 	}
 
-	// 检查每日请求限制
+	// Enforce the per-day cap after refreshing the rolling window.
 	rl.mu.RUnlock()
 	rl.mu.Lock()
 	rl.resetDailyIfNeeded(state)
@@ -260,15 +312,15 @@ func (rl *RateLimiter) IsTokenAvailable(tokenKey string) bool {
 	return true
 }
 
-// calculateBackoff 计算指数退避时间
+// calculateBackoff returns the exponential backoff for a failure count.
 func (rl *RateLimiter) calculateBackoff(failCount int) time.Duration {
 	if failCount <= 0 {
 		return 0
 	}
 
 	backoff := float64(rl.backoffBase) * math.Pow(rl.backoffMultiplier, float64(failCount-1))
 
-	// 添加抖动
+	// Add jitter so clients do not retry in lockstep.
 	jitter := backoff * rl.jitterPercent * (rl.rng.Float64()*2 - 1)
 	backoff += jitter
 
@@ -278,7 +330,7 @@ func (rl *RateLimiter) calculateBackoff(failCount int) time.Duration {
 	return time.Duration(backoff)
 }
 
-// GetTokenState 获取 Token 状态（只读）
+// GetTokenState returns a defensive copy of the token state.
 func (rl *RateLimiter) GetTokenState(tokenKey string) *TokenState {
 	rl.mu.RLock()
 	defer rl.mu.RUnlock()
@@ -288,19 +340,19 @@ func (rl *RateLimiter) GetTokenState(tokenKey string) *TokenState {
 		return nil
 	}
 
-	// 返回副本以防止外部修改
+	// Return a copy so callers cannot mutate internal state.
 	stateCopy := *state
 	return &stateCopy
 }
 
-// ClearTokenState 清除 Token 状态
+// ClearTokenState removes all tracked state for a token.
 func (rl *RateLimiter) ClearTokenState(tokenKey string) {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 	delete(rl.states, tokenKey)
 }
 
-// ResetSuspension 重置暂停状态
+// ResetSuspension clears a token's suspension and cooldown markers.
 func (rl *RateLimiter) ResetSuspension(tokenKey string) {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()