diff --git a/DevProxy.Integration.Tests/BehaviorPluginsIntegrationTests.cs b/DevProxy.Integration.Tests/BehaviorPluginsIntegrationTests.cs index bf476053..aa39cf27 100644 --- a/DevProxy.Integration.Tests/BehaviorPluginsIntegrationTests.cs +++ b/DevProxy.Integration.Tests/BehaviorPluginsIntegrationTests.cs @@ -5,6 +5,8 @@ using System.Diagnostics; using System.Globalization; using System.Net; +using System.Text; +using System.Text.Json; using DevProxy.Abstractions.Plugins; using DevProxy.Abstractions.Proxy; using DevProxy.Plugins.Behavior; @@ -212,4 +214,49 @@ public async Task RateLimitingPlusRetryAfter_ThrottledResponseCarriesRetryAfter( out _), "Retry-After header should be an integer seconds value."); } + + [Fact] + public async Task LanguageModelRateLimiting_Throttle_ReturnsOpenAIRateLimitError() + { + await using var origin = await FakeOrigin.StartAsync(); + var urls = KestrelProxyHarness.BuildUrlsToWatch(origin.Host); + + var plugin = new LanguageModelRateLimitingPlugin( + SharedHttpClient, + NullLogger.Instance, + urls, + ProxyConfig, + PluginConfig.FromJson(""" + { "promptTokenLimit": 10, "completionTokenLimit": 100, "resetTimeWindowSeconds": 300 } + """)); + + await using var proxy = await KestrelProxyHarness.StartAsync( + origin.Host, [plugin]); + using var client = proxy.CreateHttpClient(); + + // /echo returns the request body, so the usage below is read back as the + // response's token usage. #1 exceeds the prompt token limit, #2 is throttled. + const string requestBody = """ + { + "model": "gpt-4o", + "messages": [ { "role": "user", "content": "hi" } ], + "usage": { "prompt_tokens": 15, "completion_tokens": 5, "total_tokens": 20 } + } + """; + using var firstContent = new StringContent(requestBody, Encoding.UTF8, "application/json"); + using var first = await client.PostAsync(new Uri($"http://{origin.Host}/echo"), firstContent); + using var secondContent = new StringContent(requestBody, Encoding.UTF8, "application/json"); + using var throttled = await client.PostAsync(new Uri($"http://{origin.Host}/echo"), secondContent); + + Assert.Equal(HttpStatusCode.OK, first.StatusCode); + Assert.Equal(HttpStatusCode.TooManyRequests, throttled.StatusCode); + Assert.True(throttled.Headers.Contains("retry-after")); + + using var json = JsonDocument.Parse(await throttled.Content.ReadAsStringAsync()); + var error = json.RootElement.GetProperty("error"); + Assert.Equal("rate_limit_exceeded", error.GetProperty("code").GetString()); + Assert.Equal("tokens", error.GetProperty("type").GetString()); + var message = error.GetProperty("message").GetString(); + Assert.StartsWith("Rate limit reached for gpt-4o on tokens per min (TPM): Limit 10, Used 15.", message, StringComparison.Ordinal); + } } diff --git a/DevProxy.Plugins/Behavior/LanguageModelRateLimitingPlugin.cs b/DevProxy.Plugins/Behavior/LanguageModelRateLimitingPlugin.cs index a5238b1c..ae36ea37 100644 --- a/DevProxy.Plugins/Behavior/LanguageModelRateLimitingPlugin.cs +++ b/DevProxy.Plugins/Behavior/LanguageModelRateLimitingPlugin.cs @@ -51,7 +51,11 @@ public sealed class LanguageModelRateLimitingPlugin( // first request and can set the initial values private int _promptTokensRemaining = -1; private int _completionTokensRemaining = -1; + private int _promptTokensUsed; + private int _completionTokensUsed; private DateTime _resetTime = DateTime.MinValue; + // guards the token window state which is shared across concurrent connections + private readonly Lock _tokenWindowLock = new(); private LanguageModelRateLimitingCustomResponseLoader? _loader; public override string Name => nameof(LanguageModelRateLimitingPlugin); @@ -102,27 +106,10 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca return Task.CompletedTask; } - // set the initial values for the first request - if (_resetTime == DateTime.MinValue) - { - _resetTime = DateTime.Now.AddSeconds(Configuration.ResetTimeWindowSeconds); - } - if (_promptTokensRemaining == -1) - { - _promptTokensRemaining = Configuration.PromptTokenLimit; - _completionTokensRemaining = Configuration.CompletionTokenLimit; - } - - // see if we passed the reset time window - if (DateTime.Now > _resetTime) - { - _promptTokensRemaining = Configuration.PromptTokenLimit; - _completionTokensRemaining = Configuration.CompletionTokenLimit; - _resetTime = DateTime.Now.AddSeconds(Configuration.ResetTimeWindowSeconds); - } + var window = GetTokenWindowSnapshot(); // check if we have tokens available - if (_promptTokensRemaining <= 0 || _completionTokensRemaining <= 0) + if (window.PromptTokensRemaining <= 0 || window.CompletionTokensRemaining <= 0) { Logger.LogRequest($"Exceeded token limit when calling {request.Url}. Request will be throttled", MessageType.Failed, new LoggingContext(e.ProxySession)); @@ -138,9 +125,9 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca throttledRequests?.Add(new( BuildThrottleKey(request), ShouldThrottle, - _resetTime + window.ResetTime )); - ThrottleResponse(e); + ThrottleResponse(e, openAiRequest?.Model, window); state.HasBeenSet = true; } else @@ -154,7 +141,7 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca var retryAfterHeader = headersList.FirstOrDefault(h => h.Name.Equals(Configuration.HeaderRetryAfter, StringComparison.OrdinalIgnoreCase)); if (retryAfterHeader is not null && retryAfterHeader.Value == "@dynamic") { - headersList.Add(new(Configuration.HeaderRetryAfter, ((int)(_resetTime - DateTime.Now).TotalSeconds).ToString(CultureInfo.InvariantCulture))); + headersList.Add(new(Configuration.HeaderRetryAfter, ((int)(window.ResetTime - DateTime.Now).TotalSeconds).ToString(CultureInfo.InvariantCulture))); _ = headersList.Remove(retryAfterHeader); } @@ -174,7 +161,7 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca throttledRequests?.Add(new( BuildThrottleKey(request), ShouldThrottle, - _resetTime + window.ResetTime )); } @@ -194,7 +181,7 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca } else { - Logger.LogDebug("Tokens remaining - Prompt: {PromptTokensRemaining}, Completion: {CompletionTokensRemaining}", _promptTokensRemaining, _completionTokensRemaining); + Logger.LogDebug("Tokens remaining - Prompt: {PromptTokensRemaining}, Completion: {CompletionTokensRemaining}", window.PromptTokensRemaining, window.CompletionTokensRemaining); } return Task.CompletedTask; @@ -241,19 +228,28 @@ public override Task BeforeResponseAsync(ProxyResponseArgs e, CancellationToken var promptTokens = (int)openAiResponse.Usage.PromptTokens; var completionTokens = (int)openAiResponse.Usage.CompletionTokens; - _promptTokensRemaining -= promptTokens; - _completionTokensRemaining -= completionTokens; - - if (_promptTokensRemaining < 0) - { - _promptTokensRemaining = 0; - } - if (_completionTokensRemaining < 0) + int promptTokensRemaining, completionTokensRemaining; + lock (_tokenWindowLock) { - _completionTokensRemaining = 0; + _promptTokensRemaining -= promptTokens; + _completionTokensRemaining -= completionTokens; + _promptTokensUsed += promptTokens; + _completionTokensUsed += completionTokens; + + if (_promptTokensRemaining < 0) + { + _promptTokensRemaining = 0; + } + if (_completionTokensRemaining < 0) + { + _completionTokensRemaining = 0; + } + + promptTokensRemaining = _promptTokensRemaining; + completionTokensRemaining = _completionTokensRemaining; } - Logger.LogRequest($"Consumed {promptTokens} prompt tokens and {completionTokens} completion tokens. Remaining - Prompt: {_promptTokensRemaining}, Completion: {_completionTokensRemaining}", MessageType.Processed, new LoggingContext(e.ProxySession)); + Logger.LogRequest($"Consumed {promptTokens} prompt tokens and {completionTokens} completion tokens. Remaining - Prompt: {promptTokensRemaining}, Completion: {completionTokensRemaining}", MessageType.Processed, new LoggingContext(e.ProxySession)); } } catch (JsonException ex) @@ -267,34 +263,79 @@ public override Task BeforeResponseAsync(ProxyResponseArgs e, CancellationToken return Task.CompletedTask; } + private TokenWindowSnapshot GetTokenWindowSnapshot() + { + lock (_tokenWindowLock) + { + // set the initial values for the first request + if (_resetTime == DateTime.MinValue) + { + _resetTime = DateTime.Now.AddSeconds(Configuration.ResetTimeWindowSeconds); + } + if (_promptTokensRemaining == -1) + { + _promptTokensRemaining = Configuration.PromptTokenLimit; + _completionTokensRemaining = Configuration.CompletionTokenLimit; + } + + // see if we passed the reset time window + if (DateTime.Now > _resetTime) + { + _promptTokensRemaining = Configuration.PromptTokenLimit; + _completionTokensRemaining = Configuration.CompletionTokenLimit; + _promptTokensUsed = 0; + _completionTokensUsed = 0; + _resetTime = DateTime.Now.AddSeconds(Configuration.ResetTimeWindowSeconds); + } + + return new( + _resetTime, + _promptTokensRemaining, + _completionTokensRemaining, + _promptTokensUsed, + _completionTokensUsed); + } + } + private ThrottlingInfo ShouldThrottle(IHttpRequest request, string throttlingKey) { var throttleKeyForRequest = BuildThrottleKey(request); + DateTime resetTime; + lock (_tokenWindowLock) + { + resetTime = _resetTime; + } return new(throttleKeyForRequest == throttlingKey ? - (int)(_resetTime - DateTime.Now).TotalSeconds : 0, + (int)(resetTime - DateTime.Now).TotalSeconds : 0, Configuration.HeaderRetryAfter); } - private void ThrottleResponse(ProxyRequestArgs e) + private void ThrottleResponse(ProxyRequestArgs e, string? model, TokenWindowSnapshot window) { var headers = new List(); - var body = string.Empty; var request = e.ProxySession.Request; + var retryAfterSeconds = (int)(window.ResetTime - DateTime.Now).TotalSeconds; + + // Report the limit that's been exhausted, matching OpenAI's + // tokens-per-minute rate limit error so that clients back off and retry + var (limit, used) = window.PromptTokensRemaining <= 0 ? + (Configuration.PromptTokenLimit, window.PromptTokensUsed) : + (Configuration.CompletionTokenLimit, window.CompletionTokensUsed); + var modelInfo = string.IsNullOrEmpty(model) ? string.Empty : $" for {model}"; - // Build standard OpenAI error response for token limit exceeded var openAiError = new { error = new { - message = "You exceeded your current quota, please check your plan and billing details.", - type = "insufficient_quota", + message = string.Create(CultureInfo.InvariantCulture, $"Rate limit reached{modelInfo} on tokens per min (TPM): Limit {limit}, Used {used}. Please try again in {retryAfterSeconds}s."), + type = "tokens", param = (object?)null, - code = "insufficient_quota" + code = "rate_limit_exceeded" } }; - body = JsonSerializer.Serialize(openAiError, ProxyUtils.JsonSerializerOptions); + var body = JsonSerializer.Serialize(openAiError, ProxyUtils.JsonSerializerOptions); - headers.Add(new(Configuration.HeaderRetryAfter, ((int)(_resetTime - DateTime.Now).TotalSeconds).ToString(CultureInfo.InvariantCulture))); + headers.Add(new(Configuration.HeaderRetryAfter, retryAfterSeconds.ToString(CultureInfo.InvariantCulture))); if (request.Headers.Any(h => h.Name.Equals("Origin", StringComparison.OrdinalIgnoreCase))) { headers.Add(new("Access-Control-Allow-Origin", "*")); @@ -306,6 +347,13 @@ private void ThrottleResponse(ProxyRequestArgs e) private static string BuildThrottleKey(IHttpRequest r) => r.RequestUri.Host; + private readonly record struct TokenWindowSnapshot( + DateTime ResetTime, + int PromptTokensRemaining, + int CompletionTokensRemaining, + int PromptTokensUsed, + int CompletionTokensUsed); + protected override void Dispose(bool disposing) { if (disposing) diff --git a/skills/dev-proxy/references/test-llm-apps.md b/skills/dev-proxy/references/test-llm-apps.md index 8fcf7296..0fe504d1 100644 --- a/skills/dev-proxy/references/test-llm-apps.md +++ b/skills/dev-proxy/references/test-llm-apps.md @@ -166,9 +166,10 @@ Use `LanguageModelRateLimitingPlugin` to test token quota handling. ], "body": { "error": { - "message": "Token quota exceeded. Please wait.", - "type": "insufficient_quota", - "code": "token_quota_exceeded" + "message": "Rate limit reached on tokens per min (TPM). Please try again later.", + "type": "tokens", + "param": null, + "code": "rate_limit_exceeded" } } } @@ -176,6 +177,8 @@ Use `LanguageModelRateLimitingPlugin` to test token quota handling. Use `@dynamic` for the retry-after header to auto-calculate seconds until reset. +With `whenLimitExceeded: "Throttle"`, the plugin returns a 429 with OpenAI's `rate_limit_exceeded` error (`type: "tokens"`) so clients back off and retry. To simulate a billing/quota error instead (`insufficient_quota`), use `Custom` with your own response. + ### Scenario Configs **Tight limits (stress testing):** `promptTokenLimit: 500, completionTokenLimit: 250, resetTimeWindowSeconds: 30`