diff --git a/crates/jp_llm/src/error.rs b/crates/jp_llm/src/error.rs index 1e1623ba3..2605a6713 100644 --- a/crates/jp_llm/src/error.rs +++ b/crates/jp_llm/src/error.rs @@ -39,6 +39,16 @@ pub struct StreamError { /// account-wide. pub quota_scope: Option, + /// How many of the exhausted limit's usage windows the provider reported as + /// spent. + /// + /// A subscription limit reports a short window and a long one. + /// A reset credit reopens one of them, so a request stays refused while + /// more than one is spent. + /// + /// `0` when the provider reported no window state at all. + pub quota_spent_windows: usize, + /// Human-readable error message. message: String, @@ -59,6 +69,7 @@ impl StreamError { retry_after: None, quota_reset: None, quota_scope: None, + quota_spent_windows: 0, source: None, } } diff --git a/crates/jp_llm/src/provider/openai.rs b/crates/jp_llm/src/provider/openai.rs index 17e8972a4..931607d3a 100644 --- a/crates/jp_llm/src/provider/openai.rs +++ b/crates/jp_llm/src/provider/openai.rs @@ -29,6 +29,7 @@ use openai_responses::{ use reqwest::header::{self, HeaderMap, HeaderName, HeaderValue}; use serde::Deserialize; use serde_json::{Map, Value}; +use tokio::time::{Instant, sleep}; use tracing::{debug, trace, warn}; use super::{EventStream, ModelDetails, Provider}; @@ -62,6 +63,10 @@ mod fold_tests; #[path = "openai/latency_tests.rs"] mod latency_tests; +#[cfg(test)] +#[path = "openai/redeem_tests.rs"] +mod redeem_tests; + #[cfg(test)] #[path = "openai/switchable_tests.rs"] mod switchable_tests; @@ -104,6 +109,25 @@ const PERSISTED_REASONING: &str = "persisted_reasoning"; /// when the flag is present. const EXPLICIT_PROMPT_CACHING: &str = "explicit_prompt_caching"; +/// How long to keep retrying after a redeemed usage reset before giving up on +/// it. +/// +/// A redemption the account confirmed does not reach the responses host +/// instantly, so the request that prompted it is refused again for a short +/// while afterwards. +/// The credit is already spent by then, and the turn is unfinished: waiting it +/// out is what the user paid for, where abandoning the turn wastes both. +/// +/// Bounded because the alternative explanation — a window the credit does not +/// cover — looks identical from here, and that one never resolves. +const RESET_PROPAGATION_BUDGET: Duration = Duration::from_mins(1); + +/// How long to wait between attempts while a redeemed reset propagates. +/// +/// Fixed rather than backed off: the wait is already short, and every attempt +/// is an admission-stage rejection that bills nothing. +const RESET_RETRY_INTERVAL: Duration = Duration::from_secs(5); + /// How often to inject a synthetic keep-alive while a tool call is streaming. /// /// OpenAI emits the `function_call_arguments` deltas for a large tool call as a @@ -136,6 +160,15 @@ pub struct Openai { /// Where subscription requests go, after the environment override. codex_base_url: String, + /// How long a redeemed usage reset is given to reach the responses host. + /// + /// A field rather than a constant so a test can shrink the wait it is + /// pinning; production always takes [`RESET_PROPAGATION_BUDGET`]. + reset_budget: Duration, + + /// How long to wait between attempts while that reset propagates. + reset_interval: Duration, + /// The clients built for the most recently resolved credential and session. /// /// A turn issues several requests around tool execution, and resolution @@ -181,6 +214,8 @@ impl Openai { fixed_credential: None, base_url: env_override(&config.base_url_env, &config.base_url), codex_base_url: env_override(&config.codex_base_url_env, &config.codex_base_url), + reset_budget: RESET_PROPAGATION_BUDGET, + reset_interval: RESET_RETRY_INTERVAL, client_cache: Arc::new(Mutex::new(None)), }; @@ -202,10 +237,21 @@ impl Openai { fixed_credential: Some((credential, resolve::Attribution::default())), base_url: config.base_url.clone(), codex_base_url: config.codex_base_url.clone(), + reset_budget: RESET_PROPAGATION_BUDGET, + reset_interval: RESET_RETRY_INTERVAL, client_cache: Arc::new(Mutex::new(None)), } } + /// Shorten the wait for a redeemed reset, so a test pinning it runs in + /// milliseconds rather than the minute production allows. + #[cfg(test)] + pub(crate) fn with_reset_timing(mut self, budget: Duration, interval: Duration) -> Self { + self.reset_budget = budget; + self.reset_interval = interval; + self + } + /// Build a provider around an explicit subscription credential. /// /// Test seam for recording and replaying the subscription endpoint without @@ -245,11 +291,16 @@ impl Openai { /// /// Returns `None` when there is no chain to advance or nothing further in /// it, which the caller surfaces as the original, now-terminal error. + /// + /// `after_redemption` reports whether this turn already spent a reset + /// credit on the attempt's profile, which is what the recorded cooldown + /// depends on. async fn advance( &self, attempt: &resolve::Attempt, error: &StreamError, model: &str, + after_redemption: bool, ) -> Option { let spent = attempt.selected.as_ref()?; @@ -260,6 +311,7 @@ impl Openai { error, model, Utc::now(), + after_redemption, ) .await } @@ -512,7 +564,10 @@ impl Provider for Openai { // One redemption per turn. A plan holds few credits, and a window // that closes again right after being reopened is not a window a // second credit would fix. - let mut redeemed = false; + // + // The deadline is how long the reopened window is given to reach + // the responses host before the turn stops waiting on it. + let mut reopened_until: Option = None; // A credential refused or spent at admission is not a failure of // the request: the next entry in the chain can serve it. Each pass @@ -577,13 +632,18 @@ impl Provider for Openai { // subscription the user already paid for, instead of // falling through to per-token billing with allowance // still on the account. + // One credit reopens one window, so a limit reporting + // more than one spent window stays closed after the + // redemption. Spending a credit there buys nothing and + // the account has few to spend. if subscription && error.kind == StreamErrorKind::SubscriptionExhausted - && !redeemed + && reopened_until.is_none() + && error.quota_spent_windows <= 1 && let Some(notice) = this.redeem_reset_credit(&attempt, &session).await { - redeemed = true; + reopened_until = Some(Instant::now() + this.reset_budget); for notice in notices { yield Ok(Event::Notice(notice)); } @@ -591,7 +651,26 @@ impl Provider for Openai { continue; } - if let Some(mut next) = this.advance(&attempt, &error, &name).await { + // The window this turn reopened has not reached the + // host yet. The turn is unfinished and the credit is + // already spent, so it waits rather than falling + // through to per-token billing or abandoning the turn + // outright. + if subscription + && error.kind == StreamErrorKind::SubscriptionExhausted + && reopened_until.is_some_and(|until| Instant::now() < until) + { + for notice in notices { + yield Ok(Event::Notice(notice)); + } + sleep(this.reset_interval).await; + continue; + } + + if let Some(mut next) = this + .advance(&attempt, &error, &name, reopened_until.is_some()) + .await + { next.notices.splice(..0, notices); attempt = next; } else { diff --git a/crates/jp_llm/src/provider/openai/rate_limits.rs b/crates/jp_llm/src/provider/openai/rate_limits.rs index 12715a536..fac44658c 100644 --- a/crates/jp_llm/src/provider/openai/rate_limits.rs +++ b/crates/jp_llm/src/provider/openai/rate_limits.rs @@ -96,10 +96,15 @@ impl Snapshot { /// The spent window, if either is spent. #[must_use] pub fn spent(&self) -> Option<&Window> { + self.spent_windows().next() + } + + /// Every spent window of this family, in report order. + pub fn spent_windows(&self) -> impl Iterator { [self.primary.as_ref(), self.secondary.as_ref()] .into_iter() .flatten() - .find(|window| window.is_spent()) + .filter(|window| window.is_spent()) } /// The fullest window worth warning about, if any. @@ -135,6 +140,7 @@ pub fn apply(error: &mut StreamError, headers: &HeaderMap) { error.kind = StreamErrorKind::SubscriptionExhausted; error.quota_scope = Some(snapshot.scope()); error.quota_reset = window.resets_at; + error.quota_spent_windows = snapshot.spent_windows().count(); } /// The account-wide cooldown scope. diff --git a/crates/jp_llm/src/provider/openai/rate_limits_tests.rs b/crates/jp_llm/src/provider/openai/rate_limits_tests.rs index cd4754897..90d79301b 100644 --- a/crates/jp_llm/src/provider/openai/rate_limits_tests.rs +++ b/crates/jp_llm/src/provider/openai/rate_limits_tests.rs @@ -1,6 +1,7 @@ use reqwest::header::{HeaderName, HeaderValue}; use super::*; +use crate::StreamErrorKind; /// The header set a real `200` from the subscription endpoint carries. /// @@ -50,6 +51,59 @@ fn headers(pairs: &[(&str, &str)]) -> HeaderMap { map } +/// A reset credit reopens one window, so the count is what tells the caller +/// whether spending one can unblock the request at all. +#[test] +fn test_apply_counts_every_spent_window_of_the_family_it_records() { + let mut error = StreamError::other("refused"); + + apply( + &mut error, + &headers(&[ + ("x-codex-primary-used-percent", "100"), + ("x-codex-primary-window-minutes", "300"), + ("x-codex-primary-reset-after-seconds", "1800"), + ("x-codex-secondary-used-percent", "100"), + ("x-codex-secondary-window-minutes", "10080"), + ("x-codex-secondary-reset-after-seconds", "604800"), + ]), + ); + + assert_eq!(error.kind, StreamErrorKind::SubscriptionExhausted); + assert_eq!(error.quota_spent_windows, 2); +} + +#[test] +fn test_apply_counts_a_single_spent_window() { + let mut error = StreamError::other("refused"); + + apply( + &mut error, + &headers(&[ + ("x-codex-primary-used-percent", "100"), + ("x-codex-primary-window-minutes", "10080"), + ("x-codex-primary-reset-after-seconds", "604800"), + ("x-codex-secondary-used-percent", "0"), + ("x-codex-secondary-window-minutes", "0"), + ("x-codex-secondary-reset-after-seconds", "0"), + ]), + ); + + assert_eq!(error.kind, StreamErrorKind::SubscriptionExhausted); + assert_eq!(error.quota_spent_windows, 1); +} + +/// Headers reporting nothing spent leave the error alone, count included. +#[test] +fn test_apply_counts_nothing_when_no_window_is_spent() { + let mut error = StreamError::other("refused"); + + apply(&mut error, &headers(LIVE_HEADERS)); + + assert_eq!(error.kind, StreamErrorKind::Other); + assert_eq!(error.quota_spent_windows, 0); +} + #[test] fn test_parse_all_reads_both_families_from_a_live_response() { let snapshots = parse_all(&headers(LIVE_HEADERS)); diff --git a/crates/jp_llm/src/provider/openai/redeem_tests.rs b/crates/jp_llm/src/provider/openai/redeem_tests.rs new file mode 100644 index 000000000..bb6bbb547 --- /dev/null +++ b/crates/jp_llm/src/provider/openai/redeem_tests.rs @@ -0,0 +1,179 @@ +//! A redeemed usage reset keeps the turn alive while it propagates. +//! +//! The account confirms the redemption before the responses host honours it, so +//! the request that prompted it keeps being refused for a short while. +//! These tests pin what the turn does in that gap. + +use std::{ + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use futures::StreamExt as _; +use jp_config::AppConfig; +use jp_conversation::{ConversationStream, thread::Thread}; +use tokio::{ + io::{AsyncReadExt as _, AsyncWriteExt as _}, + net::TcpListener, +}; + +use super::{Openai, Provider as _}; +use crate::{model::ModelDetails, query::ChatQuery}; + +/// How many times each endpoint was called. +#[derive(Default)] +struct Calls { + responses: AtomicUsize, + consume: AtomicUsize, +} + +/// A `429` carrying the headers that mark one spent account window. +/// +/// One window, so the turn is willing to spend a credit on it; a reset an hour +/// out, so recording that timing would be visible as an hour-long cooldown. +const REFUSED: &str = concat!( + "HTTP/1.1 429 Too Many Requests\r\n", + "content-type: application/json\r\n", + "x-codex-primary-used-percent: 100\r\n", + "x-codex-primary-window-minutes: 300\r\n", + "x-codex-primary-reset-after-seconds: 3600\r\n", + "x-codex-secondary-used-percent: 0\r\n", + "x-codex-secondary-window-minutes: 0\r\n", + "x-codex-secondary-reset-after-seconds: 0\r\n", + "connection: close\r\n", + "content-length: 41\r\n\r\n", + r#"{"error":{"message":"limit reached"}}"#, +); + +fn json_response(body: &str) -> String { + format!( + "HTTP/1.1 200 OK\r\ncontent-type: application/json\r\nconnection: \ + close\r\ncontent-length: {}\r\n\r\n{body}", + body.len() + ) +} + +/// Serve the subscription host until the test drops the listener. +/// +/// Every responses request is refused, so the turn keeps waiting for the reset +/// it redeemed for as long as it is willing to. +fn spawn_host(listener: TcpListener) -> Arc { + let calls = Arc::new(Calls::default()); + let served = calls.clone(); + + tokio::spawn(async move { + loop { + let Ok((mut socket, _)) = listener.accept().await else { + return; + }; + + // Read to the end of the headers, then drain the declared body, so + // the client sees a complete exchange rather than a reset peer. + let mut request = Vec::new(); + let mut byte = [0u8; 1]; + while !request.ends_with(b"\r\n\r\n") { + match socket.read(&mut byte).await { + Ok(0) | Err(_) => break, + Ok(_) => request.push(byte[0]), + } + } + + let head = String::from_utf8_lossy(&request).to_ascii_lowercase(); + let length: usize = head + .split("content-length:") + .nth(1) + .and_then(|rest| rest.split("\r\n").next()) + .and_then(|value| value.trim().parse().ok()) + .unwrap_or(0); + let mut body = vec![0u8; length]; + drop(socket.read_exact(&mut body).await); + + let response = if head.contains("rate-limit-reset-credits/consume") { + served.consume.fetch_add(1, Ordering::SeqCst); + json_response(r#"{"success":true}"#) + } else if head.contains("rate-limit-reset-credits") { + json_response(r#"{"credits":[{"id":"credit-1"},{"id":"credit-2"}]}"#) + } else { + served.responses.fetch_add(1, Ordering::SeqCst); + REFUSED.to_owned() + }; + + drop(socket.write_all(response.as_bytes()).await); + drop(socket.shutdown().await); + } + }); + + calls +} + +fn query() -> ChatQuery { + ChatQuery::from(Thread { + system_prompt: None, + sections: vec![], + attachments: vec![], + events: ConversationStream::new_test().with_turn("hello"), + }) +} + +/// A turn that spends a credit and is refused again does not give the reset one +/// immediate retry and then abandon the turn: it keeps asking while the +/// redemption propagates. +/// +/// The bound matters as much as the retrying. +/// A window the credit does not cover is refused identically and never +/// resolves, so the wait has to end. +#[tokio::test] +async fn test_a_redeemed_reset_is_waited_out_and_then_given_up_on() { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let calls = spawn_host(listener); + + let mut config = AppConfig::new_test().providers.llm.openai; + config.codex_base_url = format!("http://{address}/codex"); + // Production allows a minute at five-second intervals. The behaviour under + // test is the shape of the wait, not its length. + let provider = Openai::with_subscription_credential( + &config, + "bearer-token".to_owned(), + "acct-1".to_owned(), + ) + .with_reset_timing(Duration::from_millis(250), Duration::from_millis(25)); + + let stream = provider + .chat_completion_stream( + &ModelDetails::empty("openai/gpt-5.6-sol".parse().unwrap()), + query(), + ) + .await + .unwrap(); + + // Far longer than the shortened budget, so it fails a wait that never ends + // rather than hanging the suite. + let events = tokio::time::timeout(Duration::from_secs(30), stream.collect::>()) + .await + .unwrap_or_else(|_| { + panic!( + "the wait for a redeemed reset never ended after {} attempts", + calls.responses.load(Ordering::SeqCst) + ) + }); + + assert!( + events.last().is_some_and(Result::is_err), + "a host that never honours the reset has to end the turn: {events:?}" + ); + + // One credit, however many times the request is refused. + assert_eq!(calls.consume.load(Ordering::SeqCst), 1); + + // Two attempts is the old behaviour: the one that hit the limit, and the + // single immediate retry after redeeming. + let attempts = calls.responses.load(Ordering::SeqCst); + assert!( + attempts > 2, + "expected the turn to keep retrying while the reset propagated, got {attempts} attempts" + ); +} diff --git a/crates/jp_llm/src/provider/openai/resolve.rs b/crates/jp_llm/src/provider/openai/resolve.rs index 2d43693ab..809b0de73 100644 --- a/crates/jp_llm/src/provider/openai/resolve.rs +++ b/crates/jp_llm/src/provider/openai/resolve.rs @@ -288,8 +288,9 @@ pub(super) async fn advance( error: &StreamError, model: &str, now: DateTime, + after_redemption: bool, ) -> Option { - record_outcome(store, spent, error, now); + record_outcome(store, spent, error, now, after_redemption); // Re-resolution reads the state just recorded against the same instant, so // a cooldown that starts now is already in effect for this walk. @@ -403,11 +404,16 @@ fn holds_refresh_token(credential: &StoredCredential, token: &str) -> bool { /// A cooldown takes a profile out of use for up to seven days and a re-login /// marker until the user acts, and a failure that says nothing about the /// credential earns neither. +/// +/// `after_redemption` says whether a reset credit was spent on this profile +/// earlier in the same turn, which changes how far the reported reset timing +/// can be trusted. fn record_outcome( store: Option<&CredentialStore>, spent: &AuthEntry, error: &StreamError, now: DateTime, + after_redemption: bool, ) { // Only a stored profile has state to record against. let AuthEntry::Subscription(Some(profile)) = spent else { @@ -424,7 +430,20 @@ fn record_outcome( } StreamErrorKind::SubscriptionExhausted | StreamErrorKind::InsufficientQuota => { let scope = error.quota_scope.as_deref().unwrap_or(SCOPE_ACCOUNT); - let until = cooldown_until(error.quota_reset, now); + + // A reset credit spent this turn reopened one of the limit's + // windows, so these headers describe a usage state JP itself just + // changed, and the window they report may not be the one the credit + // reopened. Taking their reset timing at face value can retire a + // profile for a week over a long window the user is not actually + // blocked on. The short default applies instead: it expires on its + // own, and the next request asks the provider rather than a guess. + let reported = if after_redemption { + None + } else { + error.quota_reset + }; + let until = cooldown_until(reported, now); debug!(profile, scope, %until, "Recording quota cooldown."); store.record_cooldown(CATEGORY_LLM, PROVIDER_OPENAI, profile, scope, until) } diff --git a/crates/jp_llm/src/provider/openai/resolve_tests.rs b/crates/jp_llm/src/provider/openai/resolve_tests.rs index 7a1ef8905..57478c601 100644 --- a/crates/jp_llm/src/provider/openai/resolve_tests.rs +++ b/crates/jp_llm/src/provider/openai/resolve_tests.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, sync::Arc}; use chrono::TimeZone as _; -use jp_credentials::{InMemoryCredentialBackend, MAX_COOLDOWN}; +use jp_credentials::{DEFAULT_COOLDOWN, InMemoryCredentialBackend, MAX_COOLDOWN}; use jp_storage::resource_lock::InMemoryResourceLocker; use super::*; @@ -40,6 +40,18 @@ fn token_credential(token: &str) -> StoredCredential { } } +/// The profile's account-scoped cooldown, as stored. +fn cooldown(store: &CredentialStore, profile: &str) -> Option> { + store + .load() + .unwrap() + .profiles(CATEGORY_LLM, PROVIDER_OPENAI)? + .get(profile)? + .cooldowns + .get(SCOPE_ACCOUNT) + .copied() +} + fn insert(store: &CredentialStore, profile: &str, credential: &StoredCredential) { store .mutate(|document| { @@ -416,6 +428,7 @@ async fn test_advance_without_a_selected_entry_is_terminal() { &StreamError::auth_rejected("refused"), "gpt-5.6", now(), + false, ) .await; @@ -439,6 +452,7 @@ async fn test_advance_records_relogin_and_moves_to_the_next_entry() { &StreamError::auth_rejected("token revoked"), "gpt-5.6", now(), + false, ) .await .unwrap(); @@ -479,6 +493,7 @@ async fn test_advance_records_a_cooldown_for_an_exhausted_profile() { &StreamError::new(StreamErrorKind::SubscriptionExhausted, "limit reached"), "gpt-5.6", now(), + false, ) .await .unwrap(); @@ -499,6 +514,86 @@ async fn test_advance_records_a_cooldown_for_an_exhausted_profile() { assert!(!stored.needs_relogin); } +/// Without a redemption, the reported reset is the best evidence there is, so +/// the profile stays out until the window it names reopens. +#[tokio::test] +async fn test_advance_records_the_reported_reset_for_an_exhausted_window() { + let store = store(); + insert(&store, "only", &token_credential("bearer-1")); + + let reset = now() + chrono::TimeDelta::hours(5); + let mut error = StreamError::new(StreamErrorKind::SubscriptionExhausted, "limit reached"); + error.quota_reset = Some(reset); + + advance( + &config(vec![AuthEntry::Subscription(Some("only".to_owned()))]), + Some(&store), + &AuthEntry::Subscription(Some("only".to_owned())), + &error, + "gpt-5.6", + now(), + false, + ) + .await; + + assert_eq!(cooldown(&store, "only"), Some(reset)); +} + +/// The turn already spent a reset credit against this profile, so the usage +/// state these headers describe is one JP itself just changed. +/// Recording their reset timing would retire a profile the user can still reach +/// for the full length of the window they named — a week, at the cap. +#[tokio::test] +async fn test_advance_after_a_redemption_records_only_the_short_default() { + let store = store(); + insert(&store, "only", &token_credential("bearer-1")); + + let mut error = StreamError::new(StreamErrorKind::SubscriptionExhausted, "limit reached"); + error.quota_reset = Some(now() + chrono::TimeDelta::days(30)); + + advance( + &config(vec![AuthEntry::Subscription(Some("only".to_owned()))]), + Some(&store), + &AuthEntry::Subscription(Some("only".to_owned())), + &error, + "gpt-5.6", + now(), + true, + ) + .await; + + assert_eq!(cooldown(&store, "only"), Some(now() + DEFAULT_COOLDOWN)); +} + +/// A shortened cooldown must still take the spent profile out of the walk, or +/// the chain would hand the same refused credential back. +#[tokio::test] +async fn test_advance_after_a_redemption_still_reaches_the_next_entry() { + let store = store(); + insert(&store, "first", &token_credential("bearer-1")); + insert(&store, "second", &token_credential("bearer-2")); + + let mut error = StreamError::new(StreamErrorKind::SubscriptionExhausted, "limit reached"); + error.quota_reset = Some(now() + chrono::TimeDelta::days(30)); + + let next = advance( + &config(vec![ + AuthEntry::Subscription(Some("first".to_owned())), + AuthEntry::Subscription(Some("second".to_owned())), + ]), + Some(&store), + &AuthEntry::Subscription(Some("first".to_owned())), + &error, + "gpt-5.6", + now(), + true, + ) + .await + .unwrap(); + + assert_eq!(next.credential, Credential::Bearer("bearer-2".to_owned())); +} + #[tokio::test] async fn test_advance_records_nothing_for_a_malformed_request() { // A `400` is deterministic: every credential in the chain answers it the @@ -518,6 +613,7 @@ async fn test_advance_records_nothing_for_a_malformed_request() { &StreamError::other("System messages are not allowed (HTTP 400)"), "gpt-5.6", now(), + false, ) .await; @@ -548,6 +644,7 @@ async fn test_advance_records_nothing_for_a_context_window_overflow() { &StreamError::context_window_exceeded("prompt too long"), "gpt-5.6", now(), + false, ) .await; @@ -574,6 +671,7 @@ async fn test_advance_is_terminal_when_the_chain_has_nothing_left() { &StreamError::auth_rejected("token revoked"), "gpt-5.6", now(), + false, ) .await;