Quota port: refund and settle, 429 with Retry-After and resets_at

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

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
 

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]

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 <key>` 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.

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<Vec<&'static str>>,
+}
+
+#[async_trait]
+impl Quota for RecordingQuota {
+    async fn reserve(
+        &self,
+        account: &AccountId,
+        words: usize,
+    ) -> Result<Reservation, GhostwritinError> {
+        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"]
+    );
+}

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,
         }
     }
 }