diff --git a/README.md b/README.md index d12dadc..903d4e3 100644 --- a/README.md +++ b/README.md @@ -54,7 +54,7 @@ so they can check a candidate checkout. ```sh cargo install --path crates/supermath-mcp claude mcp add supermath -- supermath-mcp serve my-project -supermath-mcp check # Lean, Python, CortexDB, Gemini key, arXiv +supermath-mcp check # Lean, Python, CortexDB, Gemini (one-token probe), arXiv ``` The project comes from the command line, never from a tool argument. A diff --git a/crates/supermath-mcp/src/check/mod.rs b/crates/supermath-mcp/src/check/mod.rs index 67f117f..37a6729 100644 --- a/crates/supermath-mcp/src/check/mod.rs +++ b/crates/supermath-mcp/src/check/mod.rs @@ -1,14 +1,16 @@ //! `supermath-mcp check`: is each backend the toolkit needs reachable? //! //! [`run`] probes the Lean toolchain and Mathlib, the Python version and -//! lock, `CortexDB`, the Gemini key, and arXiv, and returns a [`Report`]. +//! lock, `CortexDB`, Gemini, and arXiv, and returns a [`Report`]. //! A backend whose variable is unset is [`Status::Skipped`], not failed: the //! toolkit runs without it and its tools answer `BACKEND_UNAVAILABLE`. A //! configured backend that does not answer is [`Status::Failed`], and //! [`Report::passed`] is then `false`. //! -//! The Gemini key is reported as present or absent. Its value is never read -//! into the report, and nothing here prints a secret. +//! A Gemini key is probed with one ungrounded, one-token request, so a key +//! that is set but out of credit (402) or names a retired model (404) fails +//! the check. The key's value is never read into the report, and nothing +//! here prints a secret. use std::fmt; @@ -121,7 +123,9 @@ pub async fn run(toolkit: &Toolkit, config: &Config) -> Report { Line::new("cortexdb", Status::Skipped, "CORTEX_API_KEY is not set") }; let gemini = if toolkit.has_gemini_key() { - Line::new("gemini", Status::Ok, "GEMINI_API_KEY is set") + Line::from_result("gemini", toolkit.web().probe().await, |()| { + format!("{} answered", config.web.model) + }) } else { Line::new( "gemini", diff --git a/crates/supermath-mcp/src/toolkit/mod.rs b/crates/supermath-mcp/src/toolkit/mod.rs index fced03c..e3ab186 100644 --- a/crates/supermath-mcp/src/toolkit/mod.rs +++ b/crates/supermath-mcp/src/toolkit/mod.rs @@ -63,6 +63,7 @@ pub struct Toolkit { python: Python, literature: Literature, memory: Memory, + web: WebToolkit, web_has_key: bool, } @@ -191,6 +192,7 @@ impl Toolkit { python, literature, memory, + web, web_has_key: config.web.has_api_key(), }) } @@ -271,6 +273,12 @@ impl Toolkit { &self.memory } + /// The web toolkit, whose Gemini client `check` probes. + #[must_use] + pub fn web(&self) -> &WebToolkit { + &self.web + } + /// Whether a Gemini key is configured (never the key itself). #[must_use] pub fn has_gemini_key(&self) -> bool { diff --git a/crates/supermath-mcp/tests/check.rs b/crates/supermath-mcp/tests/check.rs index 1bcb20f..219e984 100644 --- a/crates/supermath-mcp/tests/check.rs +++ b/crates/supermath-mcp/tests/check.rs @@ -18,6 +18,21 @@ async fn arxiv() -> Stub { Stub::start(|_| Reply::atom(EMPTY_FEED)).await } +/// A Gemini stub answering every `generateContent` with `status`, or with a +/// one-token reply when it is 200. +async fn gemini(status: u16) -> Stub { + Stub::start(move |_| { + if status == 200 { + Reply::json(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "ok"}]}, "finishReason": "MAX_TOKENS"}] + })) + } else { + Reply::status(status) + } + }) + .await +} + async fn python() -> Stub { Stub::start(|r| match r.path.as_str() { "/health" => Reply::json( @@ -41,13 +56,15 @@ async fn every_reachable_backend_is_reported_with_what_it_found() { let python = python().await; let arxiv = arxiv().await; let cortex = cortex().await; + let gemini = gemini(200).await; let env = Env::new() .set("LEAN_URL", &lean.stub.base) .set("PYTHON_URL", &python.base) .set("SUPERMATH_ARXIV_API_URL", &arxiv.base) .set("CORTEX_API_URL", cortex.url()) .set("CORTEX_API_KEY", "c-secret") - .set("GEMINI_API_KEY", "g-secret"); + .set("GEMINI_API_KEY", "g-secret") + .set("SUPERMATH_WEB_BASE_URL", &gemini.base); let config = env.config(); let report = run(&env.toolkit(), &config).await; assert!(report.passed(), "{report}"); @@ -98,3 +115,19 @@ async fn a_configured_backend_that_does_not_answer_fails_the_check() { ); } } + +#[tokio::test] +async fn a_gemini_key_out_of_credit_fails_the_check() { + let arxiv = arxiv().await; + let gemini = gemini(402).await; + let env = Env::new() + .set("SUPERMATH_ARXIV_API_URL", &arxiv.base) + .set("GEMINI_API_KEY", "g-secret") + .set("SUPERMATH_WEB_BASE_URL", &gemini.base); + let report = run(&env.toolkit(), &env.config()).await; + assert!(!report.passed(), "{report}"); + let gemini_line = line(&report, "gemini"); + assert_eq!(gemini_line.status, Status::Failed); + assert!(gemini_line.detail.contains("credit"), "{report}"); + assert!(!report.to_string().contains("g-secret")); +} diff --git a/crates/supermath-mcp/tests/cli.rs b/crates/supermath-mcp/tests/cli.rs index c0f8a5e..836ed1d 100644 --- a/crates/supermath-mcp/tests/cli.rs +++ b/crates/supermath-mcp/tests/cli.rs @@ -50,11 +50,18 @@ fn a_bad_configuration_fails_on_stderr() { async fn check_reports_each_backend_and_exits_zero_when_none_failed() { let lean = LeanService::start().await; let arxiv = Stub::start(|_| Reply::atom(EMPTY_FEED)).await; + let gemini = Stub::start(|_| { + Reply::json(&serde_json::json!({ + "candidates": [{"content": {"parts": [{"text": "ok"}]}, "finishReason": "MAX_TOKENS"}] + })) + }) + .await; let out = binary() .args(["check", "demo"]) .env("LEAN_URL", &lean.stub.base) .env("SUPERMATH_ARXIV_API_URL", &arxiv.base) .env("GEMINI_API_KEY", "g-secret") + .env("SUPERMATH_WEB_BASE_URL", &gemini.base) .output() .unwrap(); let shown = String::from_utf8_lossy(&out.stdout); diff --git a/crates/supermath-web/README.md b/crates/supermath-web/README.md index b568316..d6051f0 100644 --- a/crates/supermath-web/README.md +++ b/crates/supermath-web/README.md @@ -36,6 +36,11 @@ the three, `register(&mut registry)` adds them to a `ToolRegistry`. - **Retired models.** Google retires model names. A 404 from Gemini is reported as `BACKEND_UNAVAILABLE` with a message naming the model and saying it may be retired; set `SUPERMATH_WEB_MODEL` to a current pinned name. +- **Out of credit.** A 402 from Gemini means the key's billing or credit is + exhausted. It is reported as `BUDGET_EXCEEDED`, not `BACKEND_UNAVAILABLE`, + so a caller does not mistake it for a transient outage and retry. + `supermath-mcp check` probes the key with one ungrounded, one-token request + (`WebToolkit::probe`), so a set but unfunded key fails the check. - **Cache and cap.** Responses are cached 24 h under `$SUPERMATH_CACHE/web`, keyed by model and prompt. `SUPERMATH_WEB_DAILY_CAP` calls per project per UTC day, counted in memory. diff --git a/crates/supermath-web/src/engine/mod.rs b/crates/supermath-web/src/engine/mod.rs index 857895d..0bb995f 100644 --- a/crates/supermath-web/src/engine/mod.rs +++ b/crates/supermath-web/src/engine/mod.rs @@ -115,6 +115,11 @@ impl Engine { } } + /// One ungrounded one-token Gemini request: does the key answer? + pub(crate) async fn probe(&self) -> Result<()> { + self.gemini.probe().await + } + /// One grounded Gemini turn: ask, read the grounding, resolve each used /// source, and keep the sentences that still have a source. pub(crate) async fn grounded_turn(&self, prompt: &str) -> Result { diff --git a/crates/supermath-web/src/error/mod.rs b/crates/supermath-web/src/error/mod.rs index e8644fe..d34ee7d 100644 --- a/crates/supermath-web/src/error/mod.rs +++ b/crates/supermath-web/src/error/mod.rs @@ -41,6 +41,13 @@ pub enum Error { /// The model name that was asked for. model: String, }, + /// Gemini answered 402: the key's billing or credit is exhausted. Not a + /// transient outage, so it reports as a spent budget, which a caller + /// does not retry. + #[error( + "gemini refused the request with 402: the key's billing or credit is exhausted, so web search is off until it is topped up" + )] + BillingExhausted, /// The upstream is down or refused the request. #[error("upstream unavailable: {0}")] Unavailable(String), @@ -109,7 +116,7 @@ impl Error { Self::Forbidden(_) => ErrorCode::Forbidden, Self::RateLimited { .. } => ErrorCode::RateLimited, Self::Timeout => ErrorCode::Timeout, - Self::DailyCapSpent { .. } => ErrorCode::BudgetExceeded, + Self::DailyCapSpent { .. } | Self::BillingExhausted => ErrorCode::BudgetExceeded, Self::InvalidConfig { .. } | Self::Client(_) | Self::MissingApiKey diff --git a/crates/supermath-web/src/error/test.rs b/crates/supermath-web/src/error/test.rs index 101c064..7bc854d 100644 --- a/crates/supermath-web/src/error/test.rs +++ b/crates/supermath-web/src/error/test.rs @@ -32,6 +32,7 @@ fn maps_every_variant_to_its_documented_code() { ), (Error::Timeout, ErrorCode::Timeout), (Error::DailyCapSpent { limit: 1 }, ErrorCode::BudgetExceeded), + (Error::BillingExhausted, ErrorCode::BudgetExceeded), (Error::InvalidUrl("x".into()), ErrorCode::InvalidInput), (Error::PrivateAddress("x".into()), ErrorCode::InvalidInput), ( diff --git a/crates/supermath-web/src/gemini/mod.rs b/crates/supermath-web/src/gemini/mod.rs index 23787b1..d736fad 100644 --- a/crates/supermath-web/src/gemini/mod.rs +++ b/crates/supermath-web/src/gemini/mod.rs @@ -102,6 +102,7 @@ impl GeminiClient { /// - [`Error::MissingApiKey`] without `GEMINI_API_KEY`. /// - [`Error::RateLimited`] on a 429, with the `Retry-After` it sent. /// - [`Error::ModelNotFound`] on a 404: the model name is likely retired. + /// - [`Error::BillingExhausted`] on a 402: the key is out of credit. /// - [`Error::Unavailable`] on any other failure status or a failed /// request. /// - [`Error::Timeout`] if the request itself times out. @@ -138,27 +139,7 @@ impl GeminiClient { .await .map_err(request_error)?; - let status = reply.status(); - if status.as_u16() == 429 { - let retry_after_ms = reply - .headers() - .get(reqwest::header::RETRY_AFTER) - .and_then(|v| v.to_str().ok()) - .and_then(|v| v.trim().parse::().ok()) - .map_or(DEFAULT_RETRY_AFTER_MS, |secs| secs.saturating_mul(1000)); - return Err(Error::RateLimited { retry_after_ms }); - } - if status.as_u16() == 404 { - return Err(Error::ModelNotFound { - model: self.model.clone(), - }); - } - if !status.is_success() { - return Err(Error::Unavailable(format!( - "gemini answered with status {}", - status.as_u16() - ))); - } + let reply = self.success(reply)?; let raw: Value = reply.json().await.map_err(|e| { if e.is_timeout() { Error::Timeout @@ -176,6 +157,62 @@ impl GeminiClient { cached_age: None, }) } + + /// Asks Gemini for one token, with no search grounding and no cache, to + /// learn whether the key and model answer: the cheapest request that + /// still meets billing. + /// + /// # Errors + /// + /// The errors of [`generate`](Self::generate) for a failure status or a + /// failed request; [`Error::BillingExhausted`] when the key is out of + /// credit. + pub async fn probe(&self) -> Result<()> { + let key = self.api_key.as_deref().ok_or(Error::MissingApiKey)?; + let url = format!( + "{}/v1beta/models/{}:generateContent", + self.base_url, self.model + ); + let body = json!({ + "contents": [{"role": "user", "parts": [{"text": "ping"}]}], + "generationConfig": {"temperature": 0, "maxOutputTokens": 1}, + }); + let reply = self + .http + .post(url) + .header("x-goog-api-key", key) + .json(&body) + .send() + .await + .map_err(request_error)?; + self.success(reply).map(|_| ()) + } + + /// `reply` itself when its status is a success, or the error its status + /// stands for. + fn success(&self, reply: reqwest::Response) -> Result { + let status = reply.status(); + match status.as_u16() { + 429 => { + let retry_after_ms = reply + .headers() + .get(reqwest::header::RETRY_AFTER) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.trim().parse::().ok()) + .map_or(DEFAULT_RETRY_AFTER_MS, |secs| secs.saturating_mul(1000)); + Err(Error::RateLimited { retry_after_ms }) + } + 404 => Err(Error::ModelNotFound { + model: self.model.clone(), + }), + 402 => Err(Error::BillingExhausted), + _ if !status.is_success() => Err(Error::Unavailable(format!( + "gemini answered with status {}", + status.as_u16() + ))), + _ => Ok(reply), + } + } } fn request_error(error: reqwest::Error) -> Error { diff --git a/crates/supermath-web/src/toolkit/mod.rs b/crates/supermath-web/src/toolkit/mod.rs index 47a982c..a2701de 100644 --- a/crates/supermath-web/src/toolkit/mod.rs +++ b/crates/supermath-web/src/toolkit/mod.rs @@ -43,6 +43,16 @@ impl WebToolkit { }) } + /// Sends Gemini the cheapest request that meets billing, to learn + /// whether the key and model answer. `check` uses it. + /// + /// # Errors + /// + /// The errors of [`GeminiClient::probe`](crate::GeminiClient::probe). + pub async fn probe(&self) -> Result<()> { + self.engine.probe().await + } + /// The `web_search` tool. #[must_use] pub fn search(&self) -> WebSearch { diff --git a/crates/supermath-web/tests/web/gemini.rs b/crates/supermath-web/tests/web/gemini.rs index cc33162..3d1efab 100644 --- a/crates/supermath-web/tests/web/gemini.rs +++ b/crates/supermath-web/tests/web/gemini.rs @@ -42,6 +42,35 @@ async fn a_404_says_the_model_may_be_retired_and_names_the_override() { assert!(message.contains("SUPERMATH_WEB_MODEL"), "{message}"); } +#[tokio::test] +async fn a_402_says_the_key_is_out_of_credit_and_is_not_a_transient_outage() { + let stub = gemini_stub(|_| Response::status(402)).await; + let client = GeminiClient::new(&config(&stub), Arc::new(SystemClock)).unwrap(); + let error = client.generate("hi").await.unwrap_err(); + assert_eq!(error, Error::BillingExhausted); + assert_eq!(error.code(), supermath_tool::ErrorCode::BudgetExceeded); + assert!(error.to_string().contains("credit"), "{error}"); +} + +#[tokio::test] +async fn a_probe_sends_one_ungrounded_token_request_and_reports_a_402() { + let stub = gemini_stub(|_| Response::status(402)).await; + let client = GeminiClient::new(&config(&stub), Arc::new(SystemClock)).unwrap(); + assert_eq!(client.probe().await.unwrap_err(), Error::BillingExhausted); + let requests = stub.requests(); + assert_eq!(requests.len(), 1); + let body = &requests[0].body; + assert!(!body.contains("google_search"), "{body}"); + assert!(body.contains("\"maxOutputTokens\":1"), "{body}"); +} + +#[tokio::test] +async fn a_probe_of_a_funded_key_succeeds() { + let stub = fixed_gemini("grounded.json").await; + let client = GeminiClient::new(&config(&stub), Arc::new(SystemClock)).unwrap(); + client.probe().await.unwrap(); +} + #[tokio::test] async fn generate_returns_the_decoded_response() { let stub = fixed_gemini("grounded.json").await;