Skip to content
Open
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
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

import com.openai.core.http.Headers;
import java.util.Objects;
import org.springframework.web.reactive.function.client.WebClientResponseException;

/**
* Resolves strict rate-limit decisions from provider headers.
Expand All @@ -12,7 +11,6 @@
*/
final class RateLimitDecisionResolver {
private static final String RETRY_AFTER_HEADER = "Retry-After";
private static final String RESET_HEADER = "X-RateLimit-Reset";

private final RateLimitHeaderParser headerParser;

Expand Down Expand Up @@ -65,31 +63,6 @@ RateLimitDecision resolveFromOpenAiRetryAfterHeaders(Headers headers) {
}
}

/**
* Resolves a decision from Spring WebClient response headers.
*
* @throws RateLimitDecisionException when headers are missing or invalid
*/
RateLimitDecision resolveFromWebClientException(WebClientResponseException webClientError) {
Objects.requireNonNull(webClientError, "webClientError");

try {
String retryAfterHeader = webClientError.getHeaders().getFirst(RETRY_AFTER_HEADER);
long retryAfterSeconds = headerParser.parseRetryAfterHeader(retryAfterHeader);
if (retryAfterHeader != null && !retryAfterHeader.isBlank()) {
return RateLimitDecision.fromRetryAfterSeconds(retryAfterSeconds);
}

return headerParser
.parseResetHeader(webClientError.getHeaders().getFirst(RESET_HEADER))
.map(RateLimitDecision::fromResetTime)
.orElseThrow(() -> new RateLimitDecisionException(
"WebClient rate-limit headers did not include Retry-After or X-RateLimit-Reset"));
} catch (IllegalArgumentException parseError) {
throw new RateLimitDecisionException("WebClient rate-limit headers are invalid", parseError);
}
}

private boolean containsHeader(Headers headers, String expectedHeaderName) {
for (String headerName : headers.names()) {
if (headerName != null && headerName.equalsIgnoreCase(expectedHeaderName)) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import org.springframework.core.env.Environment;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;
import org.springframework.web.reactive.function.client.WebClientResponseException;

/**
* Coordinates provider availability decisions from persisted and in-memory rate-limit state.
Expand Down Expand Up @@ -142,24 +141,6 @@ public void recordRateLimitFromOpenAiServiceException(ApiProvider provider, Open
applyRateLimit(requiredProvider, decision);
}

/**
* Records a rate limit from WebClient exceptions using Retry-After/X-RateLimit-Reset headers.
*
* @throws RateLimitDecisionException when headers are missing/invalid or error type is unsupported
*/
public void recordRateLimitFromException(ApiProvider provider, Throwable error) {
ApiProvider requiredProvider = Objects.requireNonNull(provider, "provider");
Throwable requiredError = Objects.requireNonNull(error, "error");

if (!(requiredError instanceof WebClientResponseException webClientError)) {
throw new RateLimitDecisionException(
"Rate-limit recording requires WebClientResponseException with headers", requiredError);
}

RateLimitDecision decision = decisionResolver.resolveFromWebClientException(webClientError);
applyRateLimit(requiredProvider, decision);
}

/**
* Clears all in-memory circuit and request counters without changing persisted state.
*/
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,11 @@
import com.openai.errors.RateLimitException;
import com.openai.errors.UnexpectedStatusCodeException;
import com.williamcallahan.javachat.support.logging.ExpectedLogEvents;
import java.nio.charset.StandardCharsets;
import java.time.Instant;
import org.junit.jupiter.api.Test;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.mock.env.MockEnvironment;
import org.springframework.web.reactive.function.client.WebClientResponseException;

/**
* Verifies strict header-only rate-limit decision behavior.
Expand Down Expand Up @@ -107,38 +104,6 @@ void recordRateLimitFromOpenAiServiceExceptionRejectsNonRateLimitStatus() {
verifyNoInteractions(rateLimitState);
}

@Test
void recordRateLimitFromExceptionUsesWebClientRetryAfterHeader() {
RateLimitState rateLimitState = mock(RateLimitState.class);
RateLimitService rateLimitService = new RateLimitService(rateLimitState, new MockEnvironment());

HttpHeaders headers = new HttpHeaders();
headers.add("Retry-After", "8");
WebClientResponseException exception = WebClientResponseException.create(
429, "Too Many Requests", headers, new byte[0], StandardCharsets.UTF_8);

try (ExpectedLogEvents expectedLogEvents = ExpectedLogEvents.capture(RATE_LIMIT_SERVICE_LOGGER)) {
rateLimitService.recordRateLimitFromException(RateLimitService.ApiProvider.OPENAI, exception);
assertRateLimitWarning(expectedLogEvents, "[openai] Rate limited (retryAfterSeconds=8)");
}

verify(rateLimitState).recordRateLimit(eq("openai"), any(Instant.class), eq("1m"));
}

@Test
void recordRateLimitFromExceptionFailsForNonWebClientErrors() {
RateLimitState rateLimitState = mock(RateLimitState.class);
RateLimitService rateLimitService = new RateLimitService(rateLimitState, new MockEnvironment());

RuntimeException exception = new RuntimeException("network issue");

assertThrows(
RateLimitDecisionException.class,
() -> rateLimitService.recordRateLimitFromException(RateLimitService.ApiProvider.OPENAI, exception));

verifyNoInteractions(rateLimitState);
}

private static void assertRateLimitWarning(ExpectedLogEvents expectedLogEvents, String expectedMessagePrefix) {
assertEquals(1, expectedLogEvents.events().size());
var rateLimitWarning = expectedLogEvents.events().getFirst();
Expand Down