Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
47 changes: 47 additions & 0 deletions DevProxy.Integration.Tests/BehaviorPluginsIntegrationTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<LanguageModelRateLimitingPlugin>.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);
}
}
134 changes: 91 additions & 43 deletions DevProxy.Plugins/Behavior/LanguageModelRateLimitingPlugin.cs
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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));

Expand All @@ -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
Expand All @@ -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);
}

Expand All @@ -174,7 +161,7 @@ public override Task BeforeRequestAsync(ProxyRequestArgs e, CancellationToken ca
throttledRequests?.Add(new(
BuildThrottleKey(request),
ShouldThrottle,
_resetTime
window.ResetTime
));
}

Expand All @@ -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;
Expand Down Expand Up @@ -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)
Expand All @@ -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<MockResponseHeader>();
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", "*"));
Expand All @@ -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)
Expand Down
9 changes: 6 additions & 3 deletions skills/dev-proxy/references/test-llm-apps.md
Original file line number Diff line number Diff line change
Expand Up @@ -166,16 +166,19 @@ 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"
}
}
}
```

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`
Expand Down
Loading