diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e696a8cb6..897befa7f4 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Added - Grok: Settings and tray **Add account** flow matching Codex/Claude — isolated `grok login --oauth`, save current CLI login, switch, and remove without logging out the active session. +- Replicate: cookie-authenticated monthly spend and optional prepaid credit balance from the billing page, with user and organization account isolation. - DeepSeek: show reported per-model spend in the provider details while preserving the billing currency, reporting period, zero values, and incomplete-total safeguards. ### Fixed diff --git a/README.md b/README.md index 577a3e317e..c610d8f17d 100755 --- a/README.md +++ b/README.md @@ -112,6 +112,7 @@ See the full history in [CHANGELOG.md](CHANGELOG.md). | Venice | API Key | USD / DIEM Balance | | OpenAI | Admin API / API Key | Usage, Requests, Project-scoped cost, Credit Balance | | Grok | Cookies / auth.json | Billing | +| Replicate | Cookies / token accounts | Monthly spend, credit balance | | ElevenLabs | API Key | Subscription Credits, Voice Slots | | Deepgram | API Key | Project Usage | | Groq | API Key | Enterprise Metrics | diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index b9b43c2f68..b7d71cf937 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/mod.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/mod.rs @@ -2,9 +2,9 @@ use std::collections::HashSet; use std::sync::Mutex; use codexbar::core::{ - FetchContext, ProviderAccountData, ProviderFetchResult, ProviderId, ProviderMetadata, - RateWindow, SourceMode, TokenAccount, TokenAccountOverride, TokenAccountStore, - instantiate_provider, + FetchContext, ManualEmptyCookiePolicy, ProviderAccountData, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, TokenAccount, TokenAccountOverride, + TokenAccountStore, instantiate_provider, }; use codexbar::locale; use codexbar::login::{self, LoginOutcome, LoginPhase}; diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs index 86c32fac88..7d26ba8fdc 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -223,6 +223,7 @@ fn cookie_source_provider(provider_id: &str) -> Option ProviderId::Sakana, "notion" => ProviderId::Notion, "grok" => ProviderId::Grok, + "replicate" => ProviderId::Replicate, _ => return None, }) } @@ -728,6 +729,22 @@ pub fn cookie_source_options_for(provider_id: &str, lang: Language) -> Vec vec![ + cookie_option( + lang, + "auto", + "Automatic imports the signed-in replicate.com browser session.", + "Paste a Cookie header from the Replicate billing page.", + None, + ), + cookie_option( + lang, + "manual", + "", + "Paste a Cookie header from https://replicate.com/account/billing.", + None, + ), + ], _ => Vec::new(), } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 7a082c4f94..0500016b22 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -85,81 +85,91 @@ pub(crate) fn build_fetch_context( let has_opencodego_api_key = id == ProviderId::OpenCodeGo && api_key.as_deref().is_some_and(|key| !key.trim().is_empty()); - let (mut source_mode, mut cookie_header) = if id.cookie_domain().is_none() { - let source_mode = if active_token_env.is_some() { - SourceMode::OAuth + let (mut source_mode, mut cookie_header, fails_closed_without_cookie) = + if id.cookie_domain().is_none() { + let source_mode = if active_token_env.is_some() { + SourceMode::OAuth + } else { + usage_source + }; + (source_mode, None, false) } else { - usage_source - }; - (source_mode, None) - } else { - match cookie_source { - // #433: an explicitly selected, non-empty Claude manual cookie is - // authoritative. Do not let an active OAuth token account silently - // replace it; this keeps tray refresh behavior aligned with diagnose, - // whose Claude Auto path tries the supplied Web cookie before OAuth. - "manual" - if provider.manual_cookie_precedes_token_account() - && stored_cookie - .as_deref() - .is_some_and(|cookie| !cookie.trim().is_empty()) => - { - (SourceMode::Web, stored_cookie.clone()) - } - _ if active_token_env.is_some() => (SourceMode::OAuth, None), - "off" if provider_uses_oauth_without_cookies(id, usage_source) => { - (SourceMode::OAuth, None) - } - "off" - if (has_kimi_code_api_key || has_opencodego_api_key) - && usage_source == SourceMode::Auto => - { - (SourceMode::Auto, None) - } - // Droid/Factory: cookie-off must never scrape browser cookies. Map to - // Cli (API-only in the provider) so Auto does not fall through to web. - "off" if id == ProviderId::Factory => (SourceMode::Cli, None), - "off" => (SourceMode::Cli, None), - "manual" => { - let cookie_header = active_token_cookie.or(stored_cookie); - let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key) - && usage_source == SourceMode::Auto + match cookie_source { + // #433: an explicitly selected, non-empty Claude manual cookie is + // authoritative. Do not let an active OAuth token account silently + // replace it; this keeps tray refresh behavior aligned with diagnose, + // whose Claude Auto path tries the supplied Web cookie before OAuth. + "manual" + if provider.manual_cookie_precedes_token_account() + && stored_cookie + .as_deref() + .is_some_and(|cookie| !cookie.trim().is_empty()) => { - SourceMode::Auto - } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source) { - // Grok Switch writes ~/.grok/auth.json. Leftover grok.com - // cookies must not force Web, or Weekly/notifications keep - // showing the previous browser account. - mode - } else if cookie_header.is_some() { - SourceMode::Web - } else if provider_uses_oauth_without_cookies(id, usage_source) { - SourceMode::OAuth - } else { - SourceMode::Cli - }; - (source_mode, cookie_header) - } - // `browser` is accepted as a legacy alias from older settings. - "auto" | "browser" | "web" => { - // Claude resolves its cached cookie and browser fallback inside - // the provider; other providers retain the shell fallback. - let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| { - if defer_provider_browser_cookie_lookup { - None + (SourceMode::Web, stored_cookie.clone(), false) + } + _ if active_token_env.is_some() => (SourceMode::OAuth, None, false), + "off" if provider_uses_oauth_without_cookies(id, usage_source) => { + (SourceMode::OAuth, None, false) + } + "off" + if (has_kimi_code_api_key || has_opencodego_api_key) + && usage_source == SourceMode::Auto => + { + (SourceMode::Auto, None, false) + } + // Droid/Factory: cookie-off must never scrape browser cookies. Map to + // Cli (API-only in the provider) so Auto does not fall through to web. + "off" if id == ProviderId::Factory => (SourceMode::Cli, None, false), + "off" => (SourceMode::Cli, None, false), + "manual" => { + let cookie_header = active_token_cookie.or(stored_cookie); + let fails_closed_without_cookie = cookie_header.is_none() + && provider.manual_empty_cookie_policy() + == ManualEmptyCookiePolicy::FailClosedWeb; + let source_mode = if (has_kimi_code_api_key || has_opencodego_api_key) + && usage_source == SourceMode::Auto + { + SourceMode::Auto + } else if let Some(mode) = grok_source_mode_for_manual_cookie(id, usage_source) + { + // Grok Switch writes ~/.grok/auth.json. Leftover grok.com + // cookies must not force Web, or Weekly/notifications keep + // showing the previous browser account. + mode + } else if cookie_header.is_some() { + SourceMode::Web + } else if fails_closed_without_cookie { + // The provider owns this policy; Web with no header means + // it fails closed instead of importing a browser account + // the user did not select. + SourceMode::Web + } else if provider_uses_oauth_without_cookies(id, usage_source) { + SourceMode::OAuth } else { - provider_cookie_domain(id, settings).and_then(|domain| { - codexbar::browser::cookies::get_cookie_header(domain) - .ok() - .filter(|h| !h.is_empty()) - }) - } - }); - (usage_source, cookie_header) + SourceMode::Cli + }; + (source_mode, cookie_header, fails_closed_without_cookie) + } + // `browser` is accepted as a legacy alias from older settings. + "auto" | "browser" | "web" => { + // Claude resolves its cached cookie and browser fallback inside + // the provider; other providers retain the shell fallback. + let cookie_header = active_token_cookie.or(stored_cookie).or_else(|| { + if defer_provider_browser_cookie_lookup { + None + } else { + provider_cookie_domain(id, settings).and_then(|domain| { + codexbar::browser::cookies::get_cookie_header(domain) + .ok() + .filter(|h| !h.is_empty()) + }) + } + }); + (usage_source, cookie_header, false) + } + _ => (usage_source, stored_cookie, false), } - _ => (usage_source, stored_cookie), - } - }; + }; // Cookie-web providers (Cursor, OpenCode, …) reject SourceMode::Cli. The shell // historically mapped "manual + no cookie" to Cli, which surfaces as @@ -207,6 +217,7 @@ pub(crate) fn build_fetch_context( FetchContext { source_mode, manual_cookie_header: cookie_header, + manual_cookie_missing: fails_closed_without_cookie, api_key, workspace_id: (!workspace_id.is_empty()).then_some(workspace_id), seat_credit_entitlement: settings.seat_credit_entitlement(id), diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index b60da4afed..8f8d7eea58 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -278,6 +278,20 @@ fn minimax_cookie_domain_follows_selected_region() { ); } +#[test] +fn replicate_cookie_source_and_domain_are_exposed() { + let mut settings = Settings::default(); + super::provider_cookie_source_set(&mut settings, "replicate", "manual".to_string()).unwrap(); + assert_eq!( + provider_cookie_source_lookup(&settings, "replicate").as_deref(), + Some("manual") + ); + assert_eq!( + super::provider_cookie_domain(ProviderId::Replicate, &settings), + Some("replicate.com") + ); +} + #[test] fn provider_cookie_source_set_rejects_unknown_provider() { let mut s = Settings::default(); @@ -455,6 +469,22 @@ fn fetch_context_opencode_empty_manual_remaps_to_web() { assert_eq!(ctx.source_mode, SourceMode::Web); } +#[test] +fn fetch_context_replicate_empty_manual_fails_closed_without_browser_import() { + let settings = Settings::default(); + let ctx = super::build_fetch_context( + ProviderId::Replicate, + &settings, + &ManualCookies::default(), + &ApiKeys::default(), + &HashMap::new(), + ); + + assert_eq!(ctx.source_mode, SourceMode::Web); + assert!(ctx.manual_cookie_header.is_none()); + assert!(ctx.manual_cookie_missing); +} + #[test] fn fetch_context_codex_manual_cookie_never_forces_unsupported_web() { // Default cookie source is "manual". Pasting a chatgpt.com cookie used to flip @@ -1710,6 +1740,13 @@ fn cookie_options_for_cookie_supporting_provider() { assert!(opts.iter().any(|o| o.label == "Disabled")); } +#[test] +fn replicate_cookie_options_allow_automatic_and_manual_sessions() { + let opts = super::cookie_source_options_for("replicate", Language::English); + let values: Vec<_> = opts.iter().map(|option| option.value.as_str()).collect(); + assert_eq!(values, vec!["auto", "manual"]); +} + #[test] fn cookie_options_empty_for_providers_without_picker() { assert!(super::cookie_source_options_for("anthropic", Language::English).is_empty()); diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-replicate.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-replicate.svg new file mode 100644 index 0000000000..6b62a2b3e9 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-replicate.svg @@ -0,0 +1 @@ + Replicate \ No newline at end of file diff --git a/apps/desktop-tauri/src/components/providers/providerIcons.ts b/apps/desktop-tauri/src/components/providers/providerIcons.ts index 58db5e5fed..2ddab175e1 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -53,6 +53,7 @@ import opencodego from "./icons/ProviderIcon-opencodego.svg?raw"; import openrouter from "./icons/ProviderIcon-openrouter.svg?raw"; import perplexity from "./icons/ProviderIcon-perplexity.svg?raw"; import qoder from "./icons/ProviderIcon-qoder.svg?raw"; +import replicate from "./icons/ProviderIcon-replicate.svg?raw"; import sakana from "./icons/ProviderIcon-sakana.svg?raw"; import stepfun from "./icons/ProviderIcon-stepfun.svg?raw"; import sub2api from "./icons/ProviderIcon-sub2api.svg?raw"; @@ -139,6 +140,7 @@ const RAW: Record = { openrouter: tint(openrouter), perplexity: tint(perplexity), qoder: tint(qoder), + replicate: tint(replicate), sakana: tint(sakana), stepfun: tint(stepfun), sub2api: tint(sub2api), @@ -216,6 +218,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { crof: { id: "crof", brandColor: "#7c3aed", fallbackLetter: "C", svgPath: RAW.crof }, crossmodel: { id: "crossmodel", brandColor: "#c084fc", fallbackLetter: "X", svgPath: RAW.crossmodel }, qoder: { id: "qoder", brandColor: "#2563eb", fallbackLetter: "Q", svgPath: RAW.qoder }, + replicate: { id: "replicate", brandColor: "#000000", fallbackLetter: "R", svgPath: RAW.replicate }, codebuddy: { id: "codebuddy", brandColor: "#0052d9", fallbackLetter: "C" }, sakana: { id: "sakana", brandColor: "#0ea5e9", fallbackLetter: "S", svgPath: RAW.sakana }, stepfun: { id: "stepfun", brandColor: "#999999", fallbackLetter: "S", svgPath: RAW.stepfun }, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index ae76cf5d5e..f752ee8d18 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -32,7 +32,7 @@ const HAS_DASHBOARD = new Set([ "azureopenai", "bedrock", "claude", "codex", "codebuff", "aiand", "commandcode", "copilot", "crof", "crossmodel", "cursor", "deepgram", "deepinfra", "deepseek", "zenmux", "clinepass", "longcat", "neuralwatt", "zoommate", "doubao", "elevenlabs", "factory", "gemini", "grok", "groq", - "infini", "jetbrains", "kilo", "kimi", "kimik2", "kiro", "manus", + "infini", "jetbrains", "kilo", "kimi", "kimik2", "kiro", "manus", "replicate", "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", "opencode", "opencodego", "openrouter", "perplexity", "qoder", "codebuddy", "sakana", "stepfun", "t3chat", "venice", "vertexai", "warp", "windsurf", diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index b281964282..380dbba0f2 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -69,6 +69,7 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["sub2api", "sub2api"], ["qwencloud", "Qwen Cloud"], ["notion", "Notion AI"], + ["replicate", "Replicate"], ["meta", "Meta"], ["muse", "Muse Code"], ]; diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index edff732cfd..237f7dab3d 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -45,6 +45,16 @@ Settings → **Providers** → provider detail → choose browser → Import. Manual cookie header paste is the fallback (required under WSL for Chromium DPAPI). Details: [COOKIES.md](./COOKIES.md). +### Replicate billing + +Replicate uses the signed-in `replicate.com` session cookie for its billing +page and read-only account endpoints. Automatic mode reuses a validated local +cookie before importing the browser session; manual mode accepts a Cookie +header containing a nonempty `sessionid`. The provider reports this month's +spend and, when the optional balance request succeeds, prepaid credit balance. +It keeps those values in the cost/detail surfaces and does not invent a quota +percentage or use a Replicate API token as a website credential. + ## Listing what is enabled ```powershell diff --git a/rust/src/browser/cookies.rs b/rust/src/browser/cookies.rs index 0d21c4251d..25be5f29dd 100755 --- a/rust/src/browser/cookies.rs +++ b/rust/src/browser/cookies.rs @@ -660,8 +660,19 @@ fn domain_matches(host_key: &str, domain: &str) -> bool { host == domain || host == format!(".{domain}") || host.ends_with(&format!(".{domain}")) } -/// Helper to get cookies for a specific domain from any available browser -pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> { +/// Per-browser cookies found while scanning every detected browser. +struct BrowserCookieCandidates { + candidates: Vec<(BrowserType, Vec)>, + abe_error_seen: bool, +} + +/// Scan every detected browser for readable cookies of `domain`. +/// +/// Returns the per-browser candidates in detection order, ignoring browsers +/// with no cookies, plus whether any browser was blocked by App-Bound +/// Encryption so callers can surface that specific, actionable error when no +/// other browser succeeded. +fn extract_domain_candidates(domain: &str) -> Result { use super::detection::BrowserDetector; let browsers = BrowserDetector::detect_all(); @@ -670,23 +681,15 @@ pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> return Err(CookieError::BrowserNotInstalled); } - // Track whether any browser raised an App-Bound Encryption error so we can - // surface that specific, actionable message if no other browser succeeds. + let mut candidates = Vec::new(); let mut abe_error_seen = false; - // Try each browser until we find cookies for browser in browsers { match CookieExtractor::extract_for_domain(&browser, domain) { Ok(cookies) if !cookies.is_empty() => { - tracing::debug!( - "Found {} cookies for {} in {}", - cookies.len(), - domain, - browser.browser_type.display_name() - ); - return Ok(cookies); + candidates.push((browser.browser_type, cookies)); } - Ok(_) => continue, + Ok(_) => {} Err(CookieError::AppBoundEncryption) => { // Chromium ABE is blocking this browser; log a warning and keep // trying the remaining browsers; Firefox does not use Chromium ABE. @@ -696,27 +699,77 @@ pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> trying remaining browsers" ); abe_error_seen = true; - // Continue to next browser rather than giving up } - Err(e) => { + Err(error) => { tracing::debug!( + browser = %browser.browser_type.display_name(), "Failed to get cookies from {}: {}", browser.browser_type.display_name(), - e + error ); } } } + Ok(BrowserCookieCandidates { + candidates, + abe_error_seen, + }) +} + +/// Helper to get cookies for a specific domain from any available browser. +/// +/// Stops at the first browser with any matching cookie; providers that must +/// try every browser after an auth failure use `get_cookie_headers_for_domain`. +pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> { + let scan = extract_domain_candidates(domain)?; + + if let Some((browser, cookies)) = scan.candidates.into_iter().next() { + tracing::debug!( + "Found {} cookies for {} in {}", + cookies.len(), + domain, + browser.display_name() + ); + return Ok(cookies); + } + // Surface a clear ABE error if it was the only kind of failure encountered, // so the UI can show an actionable message instead of a generic "not found". - if abe_error_seen { + if scan.abe_error_seen { return Err(CookieError::AppBoundEncryption); } Err(CookieError::NotFound(domain.to_string())) } +/// Get cookie-header candidates from every detected browser that has readable +/// cookies for a domain. +/// +/// Providers whose session cookie is only in one browser need the complete +/// candidate set so they can validate the session-bearing header and try the +/// next browser after an auth failure. +pub fn get_cookie_headers_for_domain( + domain: &str, +) -> Result, CookieError> { + let scan = extract_domain_candidates(domain)?; + + let headers = scan + .candidates + .into_iter() + .map(|(browser, cookies)| (browser, CookieExtractor::build_cookie_header(&cookies))) + .filter(|(_, header)| !header.trim().is_empty()) + .collect::>(); + + if headers.is_empty() && scan.abe_error_seen { + return Err(CookieError::AppBoundEncryption); + } + if headers.is_empty() { + return Err(CookieError::NotFound(domain.to_string())); + } + Ok(headers) +} + /// Get a cookie header string for a domain pub fn get_cookie_header(domain: &str) -> Result { let cookies = get_cookies_for_domain(domain)?; diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 2583604bb3..64d0b54d59 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -175,6 +175,7 @@ async fn collect_provider_diagnostic( manual_cookie_header: manual_cookies .get(provider_id.cli_name()) .map(ToOwned::to_owned), + manual_cookie_missing: false, api_key: api_keys.get(provider_id.cli_name()).map(ToOwned::to_owned), workspace_id: settings .provider_config(provider_id) diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index 7e1f5c441b..b9ecf5dce8 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -316,6 +316,7 @@ async fn fetch_guard_outcome( web_timeout, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, seat_credit_entitlement: None, diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index c7e6b350ab..2f26498cf1 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -294,6 +294,7 @@ async fn hooks_watch_observation( web_timeout, verbose, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: (!workspace.is_empty()).then(|| workspace.to_string()), seat_credit_entitlement: settings.seat_credit_entitlement(provider_id), diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 611feb6653..525d5b30f7 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -261,6 +261,10 @@ static ICONS: &[(&str, &[u8])] = &[ "ProviderIcon-qwencloud", include_bytes!("icons/ProviderIcon-qwencloud.svg"), ), + ( + "ProviderIcon-replicate", + include_bytes!("icons/ProviderIcon-replicate.svg"), + ), ( "ProviderIcon-sakana", include_bytes!("icons/ProviderIcon-sakana.svg"), diff --git a/rust/src/cli/serve/dashboard/icons/ProviderIcon-replicate.svg b/rust/src/cli/serve/dashboard/icons/ProviderIcon-replicate.svg new file mode 100644 index 0000000000..6b62a2b3e9 --- /dev/null +++ b/rust/src/cli/serve/dashboard/icons/ProviderIcon-replicate.svg @@ -0,0 +1 @@ + Replicate \ No newline at end of file diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 3c079353dd..76dc1dd162 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -156,6 +156,7 @@ async fn fetch_provider_envelope( web_timeout: 60, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, seat_credit_entitlement: None, @@ -276,6 +277,7 @@ async fn collect_claude_accounts(claude_enabled: bool) -> Option) -> String { web_timeout: 60, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, seat_credit_entitlement: None, diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index b6e8136696..f33373b90c 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -247,6 +247,7 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch web_timeout: args.web_timeout, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, seat_credit_entitlement: None, diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 2a15454183..af71e750c9 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -87,6 +87,7 @@ pub enum ProviderId { #[serde(alias = "metaspark")] Meta, Muse, + Replicate, } impl ProviderId { @@ -167,6 +168,7 @@ impl ProviderId { ProviderId::Fireworks, ProviderId::Meta, ProviderId::Muse, + ProviderId::Replicate, ] } @@ -247,6 +249,7 @@ impl ProviderId { ProviderId::QwenCloud => "qwen-cloud", ProviderId::Notion => "notion", ProviderId::Xai => "xai", + ProviderId::Replicate => "replicate", } } @@ -329,6 +332,7 @@ impl ProviderId { ProviderId::QwenCloud => "Qwen Cloud", ProviderId::Notion => "Notion AI", ProviderId::Xai => "xAI", + ProviderId::Replicate => "Replicate", } } @@ -369,6 +373,7 @@ impl ProviderId { ProviderId::CodeBuddy => Some("codebuddy.cn"), ProviderId::Sakana => Some("console.sakana.ai"), ProviderId::LongCat => Some("longcat.chat"), + ProviderId::Replicate => Some("replicate.com"), // Token-based providers (don't use cookies) ProviderId::Copilot => None, ProviderId::Zai => None, @@ -510,6 +515,7 @@ impl ProviderId { } "zoommate" | "zoom-mate" | "zoom mate" => Some(ProviderId::ZoomMate), "notion" | "notion-ai" | "notionai" | "notion ai" => Some(ProviderId::Notion), + "replicate" | "r8" => Some(ProviderId::Replicate), _ => None, } } @@ -710,6 +716,11 @@ pub struct FetchContext { /// Manual cookie header (for testing) pub manual_cookie_header: Option, + /// The cookie source is manual and no cookie is stored. The provider + /// decides what this means; Replicate fails closed instead of importing a + /// browser account the user did not select. + pub manual_cookie_missing: bool, + /// API key for providers that require authentication pub api_key: Option, @@ -746,6 +757,7 @@ impl Default for FetchContext { web_timeout: 60, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, seat_credit_entitlement: None, @@ -766,6 +778,16 @@ pub enum LastGoodFailurePolicy { PreserveOnceThenSurface, } +/// How the shell should treat a manual cookie source with no cookie present. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ManualEmptyCookiePolicy { + /// Remap to the shell's generic browser-cookie attempt. + Fallback, + /// Keep `SourceMode::Web` with no header so the provider fails closed + /// instead of importing a browser account the user did not select. + FailClosedWeb, +} + /// Trait that all providers must implement #[async_trait] pub trait Provider: Send + Sync { @@ -808,6 +830,16 @@ pub trait Provider: Send + Sync { false } + /// How the shell treats a manual cookie source with no cookie present. + /// + /// `Fallback` lets the shell remap to its generic browser-cookie attempt. + /// `FailClosedWeb` keeps `SourceMode::Web` without any header, so the + /// provider fails closed instead of importing a browser account the user + /// did not select. + fn manual_empty_cookie_policy(&self) -> ManualEmptyCookiePolicy { + ManualEmptyCookiePolicy::Fallback + } + /// Whether Automatic metric selection should prefer an exhausted quota lane. fn automatic_metric_prioritizes_exhausted_window(&self) -> bool { true @@ -1032,6 +1064,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Fireworks => "#F25B1C", ProviderId::Meta => "#0467DF", ProviderId::Muse => "#0668E1", + ProviderId::Replicate => "#000000", } } @@ -1046,7 +1079,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 74); + assert_eq!(all.len(), 75); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Fireworks)); @@ -1100,6 +1133,7 @@ mod tests { assert!(all.contains(&ProviderId::Notion)); assert!(all.contains(&ProviderId::Xai)); assert!(all.contains(&ProviderId::Meta)); + assert!(all.contains(&ProviderId::Replicate)); assert!(all.contains(&ProviderId::Muse)); } diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index a9cf6beccb..ca62d1a85a 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -18,10 +18,10 @@ use crate::providers::{ LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, - PerplexityProvider, PoeProvider, QoderProvider, QwenCloudProvider, SakanaProvider, - StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, VertexAIProvider, - WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, ZedProvider, - ZenMuxProvider, ZoomMateProvider, + PerplexityProvider, PoeProvider, QoderProvider, QwenCloudProvider, ReplicateProvider, + SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, + VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, + ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -99,6 +99,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Neuralwatt => Box::new(NeuralwattProvider::new()), ProviderId::ZoomMate => Box::new(ZoomMateProvider::new()), ProviderId::QwenCloud => Box::new(QwenCloudProvider::new()), + ProviderId::Replicate => Box::new(ReplicateProvider::new()), ProviderId::Notion => Box::new(NotionProvider::new()), ProviderId::Xai => Box::new(XaiProvider::new()), ProviderId::Fireworks => Box::new(FireworksProvider::new()), diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index b2fab33d12..fff8ae70a9 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -213,6 +213,14 @@ impl TokenAccountSupport { requires_manual_cookie_source: true, cookie_name: Some("token_v2"), }), + ProviderId::Replicate => Some(TokenAccountSupport { + title: "Session tokens", + subtitle: "Store multiple Replicate Cookie headers from the billing page.", + placeholder: "Cookie: sessionid=...; ...", + injection: TokenInjection::CookieHeader, + requires_manual_cookie_source: true, + cookie_name: Some("sessionid"), + }), ProviderId::Sub2Api => Some(TokenAccountSupport { title: "Group API keys", subtitle: "Store multiple sub2api group API keys with labels such as Claude, Codex, or Gemini.", diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 0fae2bcd84..8c7c883340 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -3,9 +3,8 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use super::ProviderDisplayDetail; use super::RateWindow; -use crate::core::ProviderDisplayDetail; - /// Subscription dates explicitly reported by an authenticated provider /// dashboard or subscription endpoint. /// diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 6d21b6ed12..e00377ac5a 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -66,6 +66,7 @@ pub mod perplexity; pub mod poe; pub mod qoder; pub mod qwencloud; +pub mod replicate; pub mod sakana; pub mod stepfun; pub mod sub2api; @@ -142,6 +143,7 @@ pub use perplexity::PerplexityProvider; pub use poe::PoeProvider; pub use qoder::QoderProvider; pub use qwencloud::QwenCloudProvider; +pub use replicate::ReplicateProvider; pub use sakana::SakanaProvider; pub use stepfun::StepFunProvider; pub use sub2api::Sub2ApiProvider; @@ -164,6 +166,19 @@ pub(crate) fn browser_cookie_header( .map_err(map_browser_cookie_error) } +pub(crate) fn browser_cookie_headers_for_domain( + domain: &str, +) -> Result, crate::core::ProviderError> { + crate::browser::cookies::get_cookie_headers_for_domain(domain) + .map(|candidates| { + candidates + .into_iter() + .map(|(browser, header)| (browser.display_name().to_string(), header)) + .collect() + }) + .map_err(map_browser_cookie_error) +} + /// All non-empty values for one cookie name in a `Cookie:` header, in order. /// Returns every value so callers can reject duplicates instead of silently /// picking the first. diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs new file mode 100644 index 0000000000..1bcebe3aa0 --- /dev/null +++ b/rust/src/providers/replicate/mod.rs @@ -0,0 +1,868 @@ +//! Replicate billing provider. +//! +//! Replicate exposes spend and prepaid credit information through its +//! authenticated billing page and read-only account endpoints. The Windows +//! port keeps credential selection native: it accepts a manually supplied +//! Cookie header or imports the `replicate.com` browser session. Browser +//! credentials stay in memory for the current fetch only. It never uses a +//! Replicate API token as a website credential and never logs cookie material. + +use async_trait::async_trait; +use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc}; +use futures::StreamExt; +use reqwest::{Client, StatusCode, Url, header::HeaderMap}; +use serde_json::Value; +use std::collections::VecDeque; +use std::time::Duration; +use tokio::time::timeout; + +use crate::core::{ + CostSnapshot, FetchContext, ManualEmptyCookiePolicy, Provider, ProviderDisplayDetail, + ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, + UsageSnapshot, +}; + +const BILLING_URL: &str = "https://replicate.com/account/billing"; +const REPLICATE_ORIGIN: &str = "https://replicate.com"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(8); +const OPTIONAL_CREDIT_TIMEOUT: Duration = Duration::from_secs(2); +const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; +const MAX_REACT_NODES: usize = 4000; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum AccountKind { + User, + Organization, +} + +impl AccountKind { + fn api_segment(self) -> &'static str { + match self { + Self::User => "users", + Self::Organization => "organizations", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct ReplicateAccount { + kind: AccountKind, + username: String, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +struct InvoiceSpend { + used: f64, +} + +pub struct ReplicateProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl ReplicateProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::Replicate, + display_name: "Replicate", + session_label: "Spend", + weekly_label: "Spend", + supports_opus: false, + supports_credits: true, + default_enabled: false, + is_primary: false, + dashboard_url: Some(BILLING_URL), + status_page_url: None, + tertiary_label_key: None, + }, + client: crate::core::credentialed_http_client_builder() + .timeout(REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_with_cookie( + &self, + cookie_header: &str, + source_label: &str, + ) -> Result { + let cookie_header = normalize_cookie_header(cookie_header).ok_or_else(|| { + ProviderError::Other( + "Replicate needs a Cookie header containing a nonempty sessionid.".to_string(), + ) + })?; + let billing_body = self + .get_text( + Url::parse(BILLING_URL).expect("valid Replicate billing URL"), + &cookie_header, + "text/html", + REQUEST_TIMEOUT, + ) + .await?; + let account = parse_billing_account(&billing_body)?; + let invoices_url = account_endpoint(&account, "invoices")?; + let invoices_body = self + .get_text( + invoices_url, + &cookie_header, + "application/json", + REQUEST_TIMEOUT, + ) + .await?; + let spend = parse_current_invoice(&invoices_body, Utc::now())?; + + let balance = self.fetch_optional_credit(&account, &cookie_header).await; + Ok(result_from_billing(account, spend, balance, source_label)) + } + + async fn fetch_optional_credit( + &self, + account: &ReplicateAccount, + cookie_header: &str, + ) -> Option { + let url = account_endpoint(account, "unused-credit").ok()?; + let body = self + .get_text( + url, + cookie_header, + "application/json", + OPTIONAL_CREDIT_TIMEOUT, + ) + .await + .ok()?; + let value: Value = serde_json::from_str(&body).ok()?; + parse_money(value.get("unused_credit")?) + } + + async fn get_text( + &self, + url: Url, + cookie_header: &str, + accept: &str, + request_timeout: Duration, + ) -> Result { + let response = timeout( + request_timeout, + self.client + .get(url) + .header("Cookie", cookie_header) + .header("Accept", accept) + .send(), + ) + .await + .map_err(|_| ProviderError::Timeout)??; + let status = response.status(); + let headers = response.headers().clone(); + validate_status(status, &headers)?; + let body = read_bounded_body(response).await?; + String::from_utf8(body).map_err(|_| { + ProviderError::Parse("Replicate returned a response that was not valid UTF-8.".into()) + }) + } + + async fn fetch_browser_cookie(&self) -> Result { + let candidates = normalized_browser_candidates( + crate::providers::browser_cookie_headers_for_domain("replicate.com")?, + ); + let mut authentication_failed = false; + for (source_label, normalized) in candidates { + match self.fetch_with_cookie(&normalized, &source_label).await { + Ok(result) => return Ok(result), + Err(error) if is_authentication_failure(&error) => { + authentication_failed = true; + } + Err(error) => return Err(error), + } + } + + if authentication_failed { + Err(ProviderError::AuthRequired) + } else { + Err(ProviderError::NoCookies) + } + } + + /// Auto and Web share one path: a manual header wins, otherwise the + /// provider tries browser candidates. There is no divergence today; if + /// Auto and Web ever need one, state it here. + async fn fetch_with_cookie_source( + &self, + ctx: &FetchContext, + ) -> Result { + if let Some(cookie_header) = ctx.manual_cookie_header.as_deref() { + return self.fetch_with_cookie(cookie_header, "manual").await; + } + // The shell signals "manual source selected, no cookie stored". Fail + // closed instead of importing a browser account the user did not + // select; browser candidates remain available for Auto without a + // manual-cookie scope. + if ctx.manual_cookie_missing { + return Err(ProviderError::Other( + "Replicate needs a Cookie header containing a nonempty sessionid.".to_string(), + )); + } + self.fetch_browser_cookie().await + } +} + +impl Default for ReplicateProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for ReplicateProvider { + fn id(&self) -> ProviderId { + ProviderId::Replicate + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::Web => self.fetch_with_cookie_source(ctx).await, + source => Err(ProviderError::UnsupportedSource(source)), + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::Web] + } + + fn supports_web(&self) -> bool { + true + } + + fn manual_cookie_precedes_token_account(&self) -> bool { + true + } + + fn manual_empty_cookie_policy(&self) -> ManualEmptyCookiePolicy { + ManualEmptyCookiePolicy::FailClosedWeb + } +} + +async fn read_bounded_body(response: reqwest::Response) -> Result, ProviderError> { + if response + .content_length() + .is_some_and(|length| length > MAX_RESPONSE_BYTES as u64) + { + return Err(response_too_large()); + } + let mut stream = response.bytes_stream(); + let mut body = Vec::new(); + while let Some(chunk) = stream.next().await { + let chunk = chunk.map_err(ProviderError::Network)?; + append_bounded_body(&mut body, &chunk)?; + } + Ok(body) +} + +fn append_bounded_body(body: &mut Vec, chunk: &[u8]) -> Result<(), ProviderError> { + if chunk.len() > MAX_RESPONSE_BYTES.saturating_sub(body.len()) { + return Err(response_too_large()); + } + body.extend_from_slice(chunk); + Ok(()) +} + +fn response_too_large() -> ProviderError { + ProviderError::Parse("Replicate returned an oversized response.".to_string()) +} + +fn account_endpoint(account: &ReplicateAccount, suffix: &str) -> Result { + let mut url = Url::parse(REPLICATE_ORIGIN) + .map_err(|_| ProviderError::Parse("Invalid Replicate endpoint.".to_string()))?; + { + let mut segments = url + .path_segments_mut() + .map_err(|_| ProviderError::Parse("Invalid Replicate endpoint.".to_string()))?; + segments + .push("api") + .push(account.kind.api_segment()) + .push(&account.username) + .push(suffix); + } + Ok(url) +} + +fn parse_billing_account(body: &str) -> Result { + let mut scanned = 0usize; + let lower = body.to_ascii_lowercase(); + let mut cursor = 0usize; + while cursor < lower.len() { + scanned += 1; + if scanned > MAX_REACT_NODES { + break; + } + let Some(relative_start) = lower[cursor..].find("') + { + cursor = after_name; + continue; + } + let Some(relative_tag_end) = lower[after_name..].find('>') else { + break; + }; + let tag_end = after_name + relative_tag_end; + let content_start = tag_end + 1; + let Some(relative_close) = lower[content_start..].find("(&body[content_start..close]) + && let Some(account) = find_account_value(&value) + { + return Ok(account); + } + cursor = close + " Option> { + let mut attrs = Vec::new(); + let mut rest = raw.trim_start(); + while !rest.is_empty() { + let name_len = rest + .chars() + .position(|c| c.is_ascii_whitespace() || c == '=') + .unwrap_or(rest.len()); + let (name, after_name) = rest.split_at(name_len); + let after_name = after_name.trim_start(); + if let Some(after_eq) = after_name.strip_prefix('=') { + let after_eq = after_eq.trim_start(); + let (value, tail) = if let Some(quoted) = after_eq.strip_prefix('"') { + quoted.split_once('"')? + } else if let Some(quoted) = after_eq.strip_prefix('\'') { + quoted.split_once('\'')? + } else { + let end = after_eq + .char_indices() + .find(|(_, c)| c.is_ascii_whitespace()) + .map(|(i, _)| i) + .unwrap_or(after_eq.len()); + after_eq.split_at(end) + }; + if !name.is_empty() { + attrs.push((name.to_ascii_lowercase(), value.to_ascii_lowercase())); + } + rest = tail.trim_start(); + } else { + if !name.is_empty() { + attrs.push((name.to_ascii_lowercase(), String::new())); + } + rest = after_name; + } + } + Some(attrs) +} + +fn find_account_value(root: &Value) -> Option { + let mut queue = VecDeque::from([root]); + let mut visited = 0usize; + while let Some(value) = queue.pop_front() { + visited += 1; + if visited > MAX_REACT_NODES { + return None; + } + if let Some(account) = value + .as_object() + .and_then(|object| object.get("account")) + .and_then(parse_account_value) + { + return Some(account); + } + match value { + Value::Array(values) => queue.extend(values), + Value::Object(object) => queue.extend(object.values()), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) => {} + } + } + None +} + +fn parse_account_value(value: &Value) -> Option { + let object = value.as_object()?; + let kind = match object.get("kind")?.as_str()?.trim() { + "user" => AccountKind::User, + "organization" => AccountKind::Organization, + _ => return None, + }; + let username = object.get("username")?.as_str()?.trim(); + if username.is_empty() || username.len() > 256 || username.chars().any(char::is_control) { + return None; + } + Some(ReplicateAccount { + kind, + username: username.to_string(), + }) +} + +fn is_signed_out_billing_page(body: &str) -> bool { + let lower = body.to_ascii_lowercase(); + let title = lower + .split_once("") + .and_then(|(_, rest)| rest.split_once("").map(|(title, _)| title.trim())) + .is_some_and(|title| title == "sign in | replicate"); + title && lower.contains("/login/github/") +} + +fn parse_current_invoice(body: &str, now: DateTime) -> Result { + let value: Value = serde_json::from_str(body).map_err(|_| parse_failure("invalid JSON"))?; + let invoices = value + .get("invoices") + .and_then(Value::as_array) + .ok_or_else(|| parse_failure("missing invoices"))?; + let current = invoices.iter().find(|invoice| { + let Some(object) = invoice.as_object() else { + return false; + }; + if object.get("type").and_then(Value::as_str) != Some("monthly-usage") { + return false; + } + match object.get("ended_before") { + None | Some(Value::Null) => true, + Some(Value::String(value)) if !value.trim().is_empty() => { + parse_invoice_end(value).is_some_and(|end| end > now) + } + _ => false, + } + }); + let current = current.ok_or_else(|| parse_failure("no current monthly-usage invoice"))?; + let used = current + .get("total_cost_before_adjustments") + .and_then(parse_money) + .ok_or_else(|| parse_failure("missing or invalid total_cost_before_adjustments"))?; + Ok(InvoiceSpend { used }) +} + +fn parse_invoice_end(value: &str) -> Option> { + let value = value.trim(); + if let Ok(end) = DateTime::parse_from_rfc3339(value) { + return Some(end.with_timezone(&Utc)); + } + if let Ok(date) = NaiveDate::parse_from_str(value, "%Y-%m-%d") { + return Some(DateTime::::from_naive_utc_and_offset( + date.and_hms_opt(0, 0, 0)?, + Utc, + )); + } + [ + "%Y-%m-%d %H:%M:%S", + "%Y-%m-%dT%H:%M:%S", + "%Y-%m-%d %H:%M:%S%.f", + "%Y-%m-%dT%H:%M:%S%.f", + ] + .into_iter() + .find_map(|format| { + NaiveDateTime::parse_from_str(value, format) + .ok() + .map(|datetime| DateTime::::from_naive_utc_and_offset(datetime, Utc)) + }) +} + +fn parse_money(value: &Value) -> Option { + let text = value.as_str()?.trim(); + if text.is_empty() + || !text + .chars() + .enumerate() + .all(|(index, character)| character.is_ascii_digit() || (character == '.' && index > 0)) + || text.matches('.').count() > 1 + || text.ends_with('.') + { + return None; + } + let number = text.parse::().ok()?; + number.is_finite().then_some(number) +} + +fn result_from_billing( + account: ReplicateAccount, + spend: InvoiceSpend, + balance: Option, + source_label: &str, +) -> ProviderFetchResult { + let spend_display = format!("${:.2}", spend.used); + let account_id = format!( + "replicate:{}:{}", + account.kind.api_segment(), + account.username + ); + let mut usage = UsageSnapshot::new(RateWindow::informational(format!( + "Spent this month: {spend_display}" + ))) + .with_login_method("Replicate"); + if account.kind == AccountKind::Organization { + usage = usage.with_organization(account.username.clone()); + } + let mut cost = CostSnapshot::new(spend.used, "USD", "This month") + .with_account_id(account.username.clone()) + .always_visible(); + if let Some(balance) = balance { + cost = cost.with_balance(balance); + } + let mut result = ProviderFetchResult::new(usage, source_label) + .with_non_authoritative_pace() + .with_cost(cost) + .with_account_identity(account_id) + .with_display_detail(ProviderDisplayDetail::new( + "spent-this-month", + "Spent this month", + spend_display, + )); + if let Some(balance) = balance { + result = result.with_display_detail(ProviderDisplayDetail::new( + "credit-balance", + "Credit balance", + format!("${balance:.2}"), + )); + } + result +} + +fn normalize_cookie_header(raw: &str) -> Option { + let mut value = raw.trim(); + if value + .get(.."cookie:".len()) + .is_some_and(|prefix| prefix.eq_ignore_ascii_case("cookie:")) + { + value = value["cookie:".len()..].trim(); + } + let mut pairs = Vec::new(); + for part in value.split(';') { + let part = part.trim(); + if part.is_empty() { + continue; + } + let (name, cookie_value) = part.split_once('=')?; + let name = name.trim(); + let cookie_value = cookie_value.trim(); + if name.is_empty() + || cookie_value.is_empty() + || name.chars().any(char::is_control) + || cookie_value.chars().any(char::is_control) + { + return None; + } + pairs.retain(|(existing, _): &(String, String)| existing != name); + pairs.push((name.to_string(), cookie_value.to_string())); + } + pairs + .iter() + .any(|(name, cookie_value)| name == "sessionid" && !cookie_value.is_empty()) + .then(|| { + pairs + .into_iter() + .map(|(name, value)| format!("{name}={value}")) + .collect::>() + .join("; ") + }) +} + +fn normalized_browser_candidates(candidates: Vec<(String, String)>) -> Vec<(String, String)> { + candidates + .into_iter() + .filter_map(|(source_label, header)| { + normalize_cookie_header(&header).map(|normalized| (source_label, normalized)) + }) + .collect() +} + +fn validate_status(status: StatusCode, headers: &HeaderMap) -> Result<(), ProviderError> { + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + let retry_after = retry_after_seconds( + headers + .get("retry-after") + .and_then(|value| value.to_str().ok()), + ); + if status == StatusCode::TOO_MANY_REQUESTS { + return Err(ProviderError::Other(format!( + "Replicate rate limit reached; retry after {retry_after:.3}s." + ))); + } + if status == StatusCode::REQUEST_TIMEOUT || status.is_server_error() { + return Err(ProviderError::Other(format!( + "Replicate billing is unavailable; retry after {retry_after:.3}s." + ))); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "Replicate returned HTTP {status}." + ))); + } + Ok(()) +} + +fn retry_after_seconds(value: Option<&str>) -> f64 { + value + .and_then(|value| value.trim().parse::().ok()) + .filter(|value| value.is_finite() && *value >= 0.0) + .map(|value| value.min(10.0)) + .unwrap_or(1.0) +} + +fn parse_failure(field: &str) -> ProviderError { + ProviderError::Parse(format!( + "Replicate billing response format changed: {field}" + )) +} + +fn is_authentication_failure(error: &ProviderError) -> bool { + matches!(error, ProviderError::AuthRequired) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn billing_page(account: &str) -> String { + format!( + r#"Billing"# + ) + } + + #[test] + fn provider_metadata_and_sources_are_cookie_only() { + let provider = ReplicateProvider::new(); + assert_eq!(provider.id(), ProviderId::Replicate); + assert_eq!(provider.metadata().display_name, "Replicate"); + assert!(!provider.metadata().default_enabled); + assert_eq!( + provider.available_sources(), + vec![SourceMode::Auto, SourceMode::Web] + ); + assert!(provider.supports_web()); + assert!(!provider.supports_cli()); + assert!(!provider.supports_oauth()); + } + + #[test] + fn parses_user_and_organization_accounts_from_bounded_react_props() { + let user = + parse_billing_account(&billing_page(r#"{"kind":"user","username":"alice"}"#)).unwrap(); + assert_eq!(user.kind, AccountKind::User); + assert_eq!(user.username, "alice"); + + let organization = parse_billing_account(&billing_page( + r#"{"kind":"organization","username":"team/acme"}"#, + )) + .unwrap(); + assert_eq!(organization.kind, AccountKind::Organization); + assert_eq!(organization.username, "team/acme"); + assert_eq!( + account_endpoint(&organization, "invoices") + .unwrap() + .as_str(), + "https://replicate.com/api/organizations/team%2Facme/invoices" + ); + } + + #[test] + fn signed_out_page_is_auth_failure_and_unknown_props_are_parse_failures() { + let signed_out = + r#" Sign in | Replicate GitHub"#; + assert!(matches!( + parse_billing_account(signed_out), + Err(ProviderError::AuthRequired) + )); + assert!(matches!( + parse_billing_account("changed"), + Err(ProviderError::Parse(_)) + )); + } + + #[test] + fn current_invoice_selection_accepts_open_and_future_invoices() { + let now = DateTime::parse_from_rfc3339("2026-09-20T00:00:00Z") + .unwrap() + .with_timezone(&Utc); + let body = serde_json::json!({ + "invoices": [ + {"type": "monthly-usage", "ended_before": "2026-09-19T00:00:00Z", "total_cost_before_adjustments": "9.00"}, + {"type": "monthly-usage", "ended_before": "2026-09-21T00:00:00Z", "total_cost_before_adjustments": "12.34"} + ] + }); + assert_eq!( + parse_current_invoice(&body.to_string(), now).unwrap().used, + 12.34 + ); + + let open = serde_json::json!({ + "invoices": [{"type": "monthly-usage", "ended_before": null, "total_cost_before_adjustments": "0"}] + }); + assert_eq!( + parse_current_invoice(&open.to_string(), now).unwrap().used, + 0.0 + ); + } + + #[test] + fn invoice_selection_accepts_date_only_and_common_naive_dates() { + let now = DateTime::parse_from_rfc3339("2026-09-20T12:00:00Z") + .unwrap() + .with_timezone(&Utc); + for ended_before in ["2026-09-21", "2026-09-21 00:00:00", "2026-09-21T00:00:00"] { + let body = serde_json::json!({ + "invoices": [{ + "type": "monthly-usage", + "ended_before": ended_before, + "total_cost_before_adjustments": "3.25" + }] + }); + assert_eq!( + parse_current_invoice(&body.to_string(), now).unwrap().used, + 3.25, + "{ended_before}" + ); + } + } + + #[test] + fn invoice_selection_fails_for_ended_missing_or_invalid_required_values() { + let now = Utc::now(); + for value in [ + serde_json::json!({"invoices": []}), + serde_json::json!({"invoices": [{"type": "monthly-usage", "ended_before": "2020-01-01T00:00:00Z", "total_cost_before_adjustments": "1"}]}), + serde_json::json!({"invoices": [{"type": "monthly-usage", "ended_before": null, "total_cost_before_adjustments": "-1"}]}), + serde_json::json!({"invoices": [{"type": "monthly-usage", "ended_before": null, "total_cost_before_adjustments": "1e2"}]}), + ] { + assert!(matches!( + parse_current_invoice(&value.to_string(), now), + Err(ProviderError::Parse(_)) + )); + } + } + + #[test] + fn money_and_optional_credit_parsing_are_strict_and_nonnegative() { + assert_eq!(parse_money(&Value::String("12.50".into())), Some(12.5)); + assert_eq!(parse_money(&Value::String("0".into())), Some(0.0)); + for text in ["", "-1", "+1", "1e2", "1.", ".5", "NaN"] { + assert_eq!(parse_money(&Value::String(text.into())), None, "{text}"); + } + let credit = serde_json::json!({"unused_credit": "4.25"}); + assert_eq!( + parse_money(credit.get("unused_credit").unwrap()), + Some(4.25) + ); + assert_eq!( + parse_money(&serde_json::json!({"unused_credit": 4.25})), + None + ); + } + + #[test] + fn result_exposes_cost_and_display_details_without_quota_math() { + let account = ReplicateAccount { + kind: AccountKind::Organization, + username: "acme".into(), + }; + let result = + result_from_billing(account, InvoiceSpend { used: 12.5 }, Some(4.25), "manual"); + assert_eq!(result.source_label, "manual"); + assert!(result.usage.primary.is_informational); + assert!(!result.pace_authoritative); + assert_eq!(result.cost.as_ref().unwrap().used, 12.5); + assert_eq!(result.cost.as_ref().unwrap().balance, Some(4.25)); + let details = result.display_details(); + assert_eq!(details.len(), 2); + assert_eq!(details[0].title(), "Spent this month"); + assert_eq!(details[1].title(), "Credit balance"); + assert_eq!(result.usage.account_organization.as_deref(), Some("acme")); + } + + #[test] + fn cookie_normalization_requires_sessionid_and_rejects_control_data() { + assert_eq!( + normalize_cookie_header("Cookie: other=1; sessionid=abc; other=2").as_deref(), + Some("sessionid=abc; other=2") + ); + assert_eq!(normalize_cookie_header("other=1"), None); + assert_eq!(normalize_cookie_header("sessionid=\r\n"), None); + } + + #[test] + fn browser_candidates_skip_headers_without_a_session_cookie() { + let candidates = normalized_browser_candidates(vec![ + ("Google Chrome".into(), "theme=dark".into()), + ("Firefox".into(), "Cookie: sessionid=valid".into()), + ]); + + assert_eq!( + candidates, + vec![("Firefox".to_string(), "sessionid=valid".to_string())] + ); + } + + #[test] + fn status_and_retry_after_classification_is_bounded() { + let headers = HeaderMap::new(); + assert!(matches!( + validate_status(StatusCode::UNAUTHORIZED, &headers), + Err(ProviderError::AuthRequired) + )); + assert!(validate_status(StatusCode::OK, &headers).is_ok()); + assert_eq!(retry_after_seconds(Some("99")), 10.0); + assert_eq!(retry_after_seconds(Some("bad")), 1.0); + assert_eq!(retry_after_seconds(Some("0.5")), 0.5); + assert!( + validate_status(StatusCode::TOO_MANY_REQUESTS, &headers) + .unwrap_err() + .to_string() + .contains("rate limit") + ); + assert!( + validate_status(StatusCode::INTERNAL_SERVER_ERROR, &headers) + .unwrap_err() + .to_string() + .contains("unavailable") + ); + + let mut retry_headers = HeaderMap::new(); + retry_headers.insert("retry-after", "0.5".parse().unwrap()); + assert!( + validate_status(StatusCode::TOO_MANY_REQUESTS, &retry_headers) + .unwrap_err() + .to_string() + .contains("0.500s") + ); + } + + #[test] + fn streamed_response_cap_rejects_oversized_chunks() { + let mut body = vec![0_u8; MAX_RESPONSE_BYTES]; + assert!(append_bounded_body(&mut body, &[0]).is_err()); + } +}