diff --git a/Cargo.lock b/Cargo.lock index 1e1c719..4dc55e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -706,6 +706,7 @@ dependencies = [ "http", "pollster", "serde_json", + "time", "tower", "tracing", ] diff --git a/Cargo.toml b/Cargo.toml index 3a6207d..4211c9a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,8 @@ cratefield-runtime-cloudflare = "0.4.2" cratefield-adapter-anthropic = "0.3.1" cratefield-adapter-openai-compatible = "0.2.1" cratefield-testing = "0.5.1" +# RFC 3339 reset times (`resets_at` on a 429). +time = { version = "0.3.45", default-features = false, features = ["std", "formatting"] } # The meaning lock and the word diff are generic, so they live in the harness # (Cratefield/harness#836, PR #837). Not released yet: a git dependency pinned # to the PR's head commit. Neither crate depends on cratefield-core, so no diff --git a/README.md b/README.md index f077e68..6690a10 100644 --- a/README.md +++ b/README.md @@ -40,7 +40,7 @@ The meaning lock and the word diff are generic, so they live in the harness, not **Works** (locally, with tests): - The engine end to end against a scripted model: voices, strengths, chunking, Markdown structure kept, the lock with one retry and then `meaning-changed`, the diff, the locks. -- The API through the harness router: auth, validation, `my_voice` without a summary (403), quotas (429), no model (503), a broken lock (422 with the locks). +- The API through the harness router: auth, validation, `my_voice` without a summary (403), quotas (429 with `Retry-After` and `resets_at`, and the words refunded when a rewrite fails), no model (503), a broken lock (422 with the locks). - The Worker builds for `wasm32-unknown-unknown` and with `worker-build`, and answered `/v1/health`, `401` and `503 model-not-configured` under `wrangler dev --local` (no model key was used, so no real rewrite went through it). A full-length rewrite also goes through `wrangler dev --local` against a scripted OpenAI-compatible server (300 ms of model latency standing in for a provider): a 9,946-word draft came back HTTP 200 in 6.5 s wall — 19 model calls, one per ~600-word chunk, strictly sequential — for about 0.2 s of workerd CPU. With a real model, wall time scales with those ~19 sequential calls; the CPU does not. - The CLI and the MCP server, against the same engine. Not run against a real provider in this repository's tests. @@ -71,7 +71,7 @@ Content-Type: application/json With a `HumanScore` wired in, a score is `{"ai_likely": 0.0–1.0, "detector": "…", "flags": [{"start", "end", "reason"}]}`. `flags` are byte offsets on char boundaries into the text that was scored (`score_before` indexes the request's text, `score_after` the rewrite), with a reason, `stock_phrase` or `uniform_rhythm`; they may overlap or nest, and they carry no words. A score without `flags` still parses. -Errors are `application/problem+json` with a stable `code`: `invalid-request`, `empty-text`, `too-many-words` (over 10,000), `unauthorized`, `voice-unavailable`, `quota-exceeded`, `meaning-changed` (with `locks`: the facts the model would not keep), `model-not-configured`, `model-unavailable`, `model-rejected`, `model-output`. The Worker adds one of its own, `rate-limited` (429, with `Retry-After`), when the account's rate-limit binding refuses the key. +Errors are `application/problem+json` with a stable `code`: `invalid-request`, `empty-text`, `too-many-words` (over 10,000), `unauthorized`, `voice-unavailable`, `quota-exceeded` (with `resets_at`, RFC 3339, and a `Retry-After` header), `meaning-changed` (with `locks`: the facts the model would not keep), `model-not-configured`, `model-unavailable`, `model-rejected`, `model-output`. The Worker adds one of its own, `rate-limited` (429, with `Retry-After`), when the account's rate-limit binding refuses the key. ### The meaning lock diff --git a/crates/ghostwritin-api/Cargo.toml b/crates/ghostwritin-api/Cargo.toml index dbfa28b..129ad79 100644 --- a/crates/ghostwritin-api/Cargo.toml +++ b/crates/ghostwritin-api/Cargo.toml @@ -19,6 +19,7 @@ axum = { workspace = true } bytes = { workspace = true } http = { workspace = true } serde_json = { workspace = true } +time = { workspace = true } tracing = { workspace = true } [dev-dependencies] diff --git a/crates/ghostwritin-api/src/lib.rs b/crates/ghostwritin-api/src/lib.rs index a114f0e..f1a52a9 100644 --- a/crates/ghostwritin-api/src/lib.rs +++ b/crates/ghostwritin-api/src/lib.rs @@ -12,7 +12,10 @@ //! repository's: hashed static API keys, no quota, no human score, no //! watermark detection, no My voice. The hosted service builds the same //! modules with its own `Services` in its own Worker crate; nothing here -//! needs to change for that. +//! needs to change for that. Plans and billing are hosted-only: the +//! hosted service brings its own `Quota` (reserve, refund, settle) and its +//! own billing module, and reuses [`bearer`] and [`problem`] so its routes +//! answer like these. //! //! # Data policy //! @@ -249,16 +252,31 @@ async fn handle( ), _ => None, }; - services.quota.reserve(&account, words).await?; + // The words are charged before the model is called and refunded if it + // fails; a rewrite that succeeded settles them (where a plan meters). + let reservation = services.quota.reserve(&account, words).await?; - Engine::new(model) + match Engine::new(model) .human_score(services.human_score.clone()) .watermarks(services.watermarks.clone()) .rewrite(&request, style.as_ref()) .await + { + Ok(answer) => { + services.quota.settle(&reservation).await; + Ok(answer) + } + Err(error) => { + if let Err(refund) = services.quota.refund(&reservation).await { + tracing::warn!(target: "ghostwritin::rewrite", code = refund.code(), "quota refund failed"); + } + Err(error) + } + } } -fn bearer(headers: &HeaderMap) -> Option<&str> { +/// The key from an `Authorization: Bearer ` header, if there is one. +pub fn bearer(headers: &HeaderMap) -> Option<&str> { let value = headers.get(header::AUTHORIZATION)?.to_str().ok()?; let (scheme, key) = value.split_once(' ')?; let key = key.trim(); @@ -267,7 +285,9 @@ fn bearer(headers: &HeaderMap) -> Option<&str> { /// An RFC 9457 problem. `locks` rides along for `meaning-changed`, because /// the caller needs to know which facts broke; it goes to the caller only. -fn problem(error: &GhostwritinError) -> Response { +/// A used-up quota says when it comes back: `resets_at` (RFC 3339) in the +/// body, `Retry-After` in the headers. +pub fn problem(error: &GhostwritinError) -> Response { let status = StatusCode::from_u16(error.status()).unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); let mut body = json!({ "type": format!("{PROBLEM_BASE}{}", error.code()), @@ -278,11 +298,21 @@ fn problem(error: &GhostwritinError) -> Response { if let GhostwritinError::MeaningChanged { locks } = error { body["locks"] = serde_json::to_value(locks).unwrap_or_default(); } + if let GhostwritinError::QuotaExceeded { resets_at, .. } = error { + body["resets_at"] = json!(rfc3339(*resets_at)); + } let mut response = (status, axum::Json(body)).into_response(); response.headers_mut().insert( header::CONTENT_TYPE, HeaderValue::from_static("application/problem+json"), ); + if let GhostwritinError::QuotaExceeded { retry_after, .. } = error { + response.headers_mut().insert( + header::RETRY_AFTER, + HeaderValue::from_str(&retry_after.to_string()) + .unwrap_or(HeaderValue::from_static("0")), + ); + } if matches!(error, GhostwritinError::Unauthorized) { response .headers_mut() @@ -292,6 +322,19 @@ fn problem(error: &GhostwritinError) -> Response { response } +/// Unix seconds as `YYYY-MM-DDTHH:MM:SSZ` (RFC 3339, UTC): the wire form +/// of a quota's reset time. +fn rfc3339(unix_secs: u64) -> String { + use time::format_description::well_known::Rfc3339; + let at = + time::OffsetDateTime::from_unix_timestamp(i64::try_from(unix_secs).unwrap_or(i64::MAX)) + .unwrap_or(time::OffsetDateTime::UNIX_EPOCH); + at.replace_nanosecond(0) + .unwrap_or(at) + .format(&Rfc3339) + .unwrap_or_default() +} + /// [`RequestLog`]s as `tracing` events (target `ghostwritin::rewrite`), for /// native runtimes and tests. On Workers `tracing` has no subscriber (the /// harness installs none on wasm), so the Worker uses a console logger. diff --git a/crates/ghostwritin-api/tests/routes.rs b/crates/ghostwritin-api/tests/routes.rs index cd8bb18..43a27b4 100644 --- a/crates/ghostwritin-api/tests/routes.rs +++ b/crates/ghostwritin-api/tests/routes.rs @@ -8,12 +8,13 @@ use axum::body::{Body, to_bytes}; use cratefield_core::{Completion, Prompt, TextModel, TextModelError}; use cratefield_testing::{TestHarness, conformance}; use ghostwritin_api::{HealthApi, RewriteApi, Services}; -use ghostwritin_core::AccountId; use ghostwritin_core::ports::{ - DailyWordLimit, InMemoryVoiceStore, RequestLog, RequestLogger, StaticApiKeys, VoiceStore, + DailyWordLimit, InMemoryVoiceStore, Quota, RequestLog, RequestLogger, Reservation, + StaticApiKeys, VoiceStore, }; +use ghostwritin_core::{AccountId, GhostwritinError}; use ghostwritin_core::{Strength, Voice}; -use http::{Method, Request, StatusCode, header}; +use http::{HeaderMap, Method, Request, StatusCode, header}; use serde_json::{Value, json}; use tower::ServiceExt; @@ -82,6 +83,18 @@ async fn call( key: Option<&str>, body: &str, ) -> (StatusCode, Value, String) { + let (status, json, content_type, _) = call_with_headers(kit, method, path, key, body).await; + (status, json, content_type) +} + +/// [`call`], with the response headers (`Retry-After` on a used-up quota). +async fn call_with_headers( + kit: &TestHarness, + method: Method, + path: &str, + key: Option<&str>, + body: &str, +) -> (StatusCode, Value, String, HeaderMap) { let mut request = Request::builder() .method(method) .uri(path) @@ -96,6 +109,7 @@ async fn call( .await .unwrap(); let status = response.status(); + let response_headers = response.headers().clone(); let content_type = response .headers() .get(header::CONTENT_TYPE) @@ -106,6 +120,7 @@ async fn call( status, serde_json::from_slice(&bytes).unwrap_or(Value::Null), content_type, + response_headers, ) } @@ -291,7 +306,7 @@ async fn my_voice_uses_the_stored_summary_only() { #[pollster::test] async fn a_quota_is_429() { let mut limited = services(); - limited.quota = Arc::new(DailyWordLimit::new(30, || 1)); + limited.quota = Arc::new(DailyWordLimit::new(30, || 86_400)); let kit = kit_with(limited, Some(polish)); let (status, _, _) = call( &kit, @@ -302,7 +317,7 @@ async fn a_quota_is_429() { ) .await; assert_eq!(status, StatusCode::OK); - let (status, json, _) = call( + let (status, json, content_type, headers) = call_with_headers( &kit, Method::POST, "/v1/rewrite", @@ -312,6 +327,9 @@ async fn a_quota_is_429() { .await; assert_eq!(status, StatusCode::TOO_MANY_REQUESTS); assert_eq!(json["code"], "quota-exceeded"); + assert_eq!(json["resets_at"], "1970-01-03T00:00:00Z"); + assert_eq!(headers.get(header::RETRY_AFTER).unwrap(), "86400"); + assert_eq!(content_type, "application/problem+json"); } #[derive(Default)] @@ -370,3 +388,71 @@ fn modules_conform() { conformance(Box::new(RewriteApi::new(services()))); conformance(Box::new(HealthApi)); } + +/// A quota that records what the API did with each reservation, over an +/// in-memory daily limit. +struct RecordingQuota { + inner: DailyWordLimit, + calls: Mutex>, +} + +#[async_trait] +impl Quota for RecordingQuota { + async fn reserve( + &self, + account: &AccountId, + words: usize, + ) -> Result { + self.calls.lock().unwrap().push("reserve"); + self.inner.reserve(account, words).await + } + + async fn refund(&self, reservation: &Reservation) -> Result<(), GhostwritinError> { + self.calls.lock().unwrap().push("refund"); + self.inner.refund(reservation).await + } + + async fn settle(&self, _reservation: &Reservation) { + self.calls.lock().unwrap().push("settle"); + } +} + +#[pollster::test] +async fn a_failed_rewrite_refunds_and_a_good_one_settles() { + let quota = Arc::new(RecordingQuota { + // Room for one draft a day, not two. + inner: DailyWordLimit::new(30, || 86_400), + calls: Mutex::new(Vec::new()), + }); + let mut limited = services(); + limited.quota = quota.clone(); + + // The model breaks a locked number: 422, and the words go back. + let kit = kit_with(limited.clone(), Some(break_numbers)); + let (status, json, _) = call( + &kit, + Method::POST, + "/v1/rewrite", + Some("test-key"), + &body("casual"), + ) + .await; + assert_eq!(status, StatusCode::UNPROCESSABLE_ENTITY, "{json}"); + assert_eq!(*quota.calls.lock().unwrap(), ["reserve", "refund"]); + + // So the same draft still fits, and a good rewrite settles its words. + let kit = kit_with(limited, Some(polish)); + let (status, json, _) = call( + &kit, + Method::POST, + "/v1/rewrite", + Some("test-key"), + &body("casual"), + ) + .await; + assert_eq!(status, StatusCode::OK, "{json}"); + assert_eq!( + *quota.calls.lock().unwrap(), + ["reserve", "refund", "reserve", "settle"] + ); +} diff --git a/crates/ghostwritin-core/src/error.rs b/crates/ghostwritin-core/src/error.rs index ec855a0..d3a4e7d 100644 --- a/crates/ghostwritin-core/src/error.rs +++ b/crates/ghostwritin-core/src/error.rs @@ -38,9 +38,19 @@ pub enum GhostwritinError { #[error("the language model's answer did not have the expected shape")] ModelOutput, #[error("the word quota for this account is used up")] - QuotaExceeded, + QuotaExceeded { + /// Unix seconds when the account's period rolls over and the words + /// come back. + resets_at: u64, + /// Seconds from now until `resets_at`: the `Retry-After` value. + retry_after: u64, + }, #[error("a valid API key is required")] Unauthorized, + #[error("billing is not configured")] + BillingNotConfigured, + #[error("billing is unavailable; try again shortly")] + BillingUnavailable, #[error("storage failed")] Storage, } @@ -60,8 +70,10 @@ impl GhostwritinError { Self::ModelUnavailable => "model-unavailable", Self::ModelRejected => "model-rejected", Self::ModelOutput => "model-output", - Self::QuotaExceeded => "quota-exceeded", + Self::QuotaExceeded { .. } => "quota-exceeded", Self::Unauthorized => "unauthorized", + Self::BillingNotConfigured => "billing-not-configured", + Self::BillingUnavailable => "billing-unavailable", Self::Storage => "storage", } } @@ -77,9 +89,13 @@ impl GhostwritinError { Self::Unauthorized => 401, Self::VoiceUnavailable => 403, Self::MeaningChanged { .. } => 422, - Self::QuotaExceeded => 429, + Self::QuotaExceeded { .. } => 429, Self::ModelRejected | Self::ModelOutput => 502, - Self::ModelNotConfigured | Self::ModelUnavailable | Self::Storage => 503, + Self::ModelNotConfigured + | Self::ModelUnavailable + | Self::BillingNotConfigured + | Self::BillingUnavailable + | Self::Storage => 503, } } } diff --git a/crates/ghostwritin-core/src/ports.rs b/crates/ghostwritin-core/src/ports.rs index 1f72fc0..d413f61 100644 --- a/crates/ghostwritin-core/src/ports.rs +++ b/crates/ghostwritin-core/src/ports.rs @@ -75,16 +75,55 @@ impl WatermarkDetector for NoWatermarkDetector { } } +/// The seconds in a UTC day, the window [`DailyWordLimit`] counts in. +const SECS_PER_DAY: u64 = 86_400; + +/// Words held for one rewrite: what [`Quota::reserve`] charged and where, +/// so [`Quota::refund`] can put them back into the same window after a +/// failed rewrite. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Reservation { + pub account: AccountId, + pub words: usize, + /// The window the words went to, as the implementation keys it + /// (a day number, a month index, whatever the quota counts in). + pub period: u64, + /// Unix seconds when that window rolls over. + pub resets_at: u64, +} + /// Word quotas: reserve the words a rewrite will use before the model is -/// called. +/// called, refund them if it fails, settle them once it succeeded. #[async_trait] pub trait Quota: Send + Sync { /// # Errors /// - /// [`GhostwritinError::QuotaExceeded`] when `words` more would go over - /// the account's allowance; [`GhostwritinError::Storage`] when the - /// count cannot be read. - async fn reserve(&self, account: &AccountId, words: usize) -> Result<(), GhostwritinError>; + /// [`GhostwritinError::QuotaExceeded`] (carrying the period's reset + /// time) when `words` more would go over the account's allowance; + /// [`GhostwritinError::Storage`] when the count cannot be read or + /// written. + async fn reserve( + &self, + account: &AccountId, + words: usize, + ) -> Result; + + /// Puts a failed rewrite's words back. A refund after the period + /// rolled over, or a refund applied twice, must not over-credit: + /// implementations count the same window the reservation named and + /// never below zero. + /// + /// # Errors + /// + /// [`GhostwritinError::Storage`] when the count cannot be written. + async fn refund(&self, reservation: &Reservation) -> Result<(), GhostwritinError>; + + /// The rewrite the words were reserved for succeeded: they are spent. + /// A hook for quotas that bill what was used (a metered plan reports + /// the count here); the open ones have nothing to do. Infallible on + /// purpose: settling must never fail a rewrite that already succeeded, + /// so an implementation logs its own failures (counts only, no text). + async fn settle(&self, _reservation: &Reservation) {} } /// No quota: every reservation succeeds. The CLI's and the MCP server's, @@ -94,7 +133,20 @@ pub struct Unlimited; #[async_trait] impl Quota for Unlimited { - async fn reserve(&self, _account: &AccountId, _words: usize) -> Result<(), GhostwritinError> { + async fn reserve( + &self, + account: &AccountId, + words: usize, + ) -> Result { + Ok(Reservation { + account: account.clone(), + words, + period: 0, + resets_at: 0, + }) + } + + async fn refund(&self, _reservation: &Reservation) -> Result<(), GhostwritinError> { Ok(()) } } @@ -103,38 +155,67 @@ impl Quota for Unlimited { /// /// Basic on purpose: the count lives in this process, so it resets on a /// restart and is per isolate on Workers. Good enough for one self-hosted -/// server; the hosted service keeps its counts in a database. The day is -/// whatever `today` returns (days since an epoch), so tests and runtimes -/// without a system clock can supply it. +/// server; the hosted service keeps its counts in a database. `now` is +/// unix seconds, so +/// tests and runtimes without a system clock can supply it; the day is +/// derived from it. pub struct DailyWordLimit { limit: usize, - today: Arc u64 + Send + Sync>, + now: Arc u64 + Send + Sync>, used: Mutex>, } impl DailyWordLimit { - pub fn new(limit: usize, today: impl Fn() -> u64 + Send + Sync + 'static) -> Self { + pub fn new(limit: usize, now: impl Fn() -> u64 + Send + Sync + 'static) -> Self { Self { limit, - today: Arc::new(today), + now: Arc::new(now), used: Mutex::new(HashMap::new()), } } + + /// The window `now` (unix seconds) falls in. + fn day(now: u64) -> (u64, u64) { + (now / SECS_PER_DAY, (now / SECS_PER_DAY + 1) * SECS_PER_DAY) + } } #[async_trait] impl Quota for DailyWordLimit { - async fn reserve(&self, account: &AccountId, words: usize) -> Result<(), GhostwritinError> { - let today = (self.today)(); + async fn reserve( + &self, + account: &AccountId, + words: usize, + ) -> Result { + let now = (self.now)(); + let (day, resets_at) = Self::day(now); let mut used = self.used.lock().map_err(|_| GhostwritinError::Storage)?; - let entry = used.entry(account.clone()).or_insert((today, 0)); - if entry.0 != today { - *entry = (today, 0); + let entry = used.entry(account.clone()).or_insert((day, 0)); + if entry.0 != day { + *entry = (day, 0); } if entry.1.saturating_add(words) > self.limit { - return Err(GhostwritinError::QuotaExceeded); + return Err(GhostwritinError::QuotaExceeded { + resets_at, + retry_after: resets_at.saturating_sub(now), + }); } entry.1 += words; + Ok(Reservation { + account: account.clone(), + words, + period: day, + resets_at, + }) + } + + async fn refund(&self, reservation: &Reservation) -> Result<(), GhostwritinError> { + let mut used = self.used.lock().map_err(|_| GhostwritinError::Storage)?; + if let Some(entry) = used.get_mut(&reservation.account) + && entry.0 == reservation.period + { + entry.1 = entry.1.saturating_sub(reservation.words); + } Ok(()) } } @@ -347,20 +428,41 @@ mod tests { #[pollster::test] async fn a_daily_limit_resets_each_day() { - let day = Arc::new(std::sync::atomic::AtomicU64::new(1)); + let day = Arc::new(std::sync::atomic::AtomicU64::new(86_400)); let clock = day.clone(); let quota = DailyWordLimit::new(100, move || clock.load(std::sync::atomic::Ordering::SeqCst)); let me = account("me"); - assert_eq!(quota.reserve(&me, 60).await, Ok(())); + let first = quota.reserve(&me, 60).await.unwrap(); + assert_eq!(first.period, 1); + assert_eq!(first.resets_at, 172_800); assert_eq!( quota.reserve(&me, 41).await, - Err(GhostwritinError::QuotaExceeded) + Err(GhostwritinError::QuotaExceeded { + resets_at: 172_800, + retry_after: 86_400, + }) ); - assert_eq!(quota.reserve(&account("other"), 100).await, Ok(())); - assert_eq!(quota.reserve(&me, 40).await, Ok(())); - day.store(2, std::sync::atomic::Ordering::SeqCst); - assert_eq!(quota.reserve(&me, 100).await, Ok(())); + assert!(quota.reserve(&account("other"), 100).await.is_ok()); + // A failed rewrite is refunded, and the words are back. + assert_eq!(quota.refund(&first).await, Ok(())); + assert!(quota.reserve(&me, 60).await.is_ok()); + // The next day the whole allowance is back. + day.store(172_800, std::sync::atomic::Ordering::SeqCst); + assert!(quota.reserve(&me, 100).await.is_ok()); + // A refund for yesterday's window does not credit today's. + assert_eq!(quota.refund(&first).await, Ok(())); + assert!(quota.reserve(&me, 1).await.is_err()); + } + + #[pollster::test] + async fn unlimited_reserves_and_refunds_without_counting() { + let quota = Unlimited; + let me = account("me"); + let reservation = quota.reserve(&me, 10_000).await.unwrap(); + assert_eq!(reservation.words, 10_000); + assert_eq!(quota.refund(&reservation).await, Ok(())); + quota.settle(&reservation).await; } #[pollster::test]