From 4ceb88f03c76177e56a9459997999877a364d231 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 19 Sep 2026 21:34:59 +0700 Subject: [PATCH 01/12] Add transient provider inventory plumbing --- .../src-tauri/src/commands/bridge.rs | 24 +++++ .../src-tauri/src/commands/provider_detail.rs | 3 + .../src-tauri/src/commands/providers.rs | 1 + .../src-tauri/src/commands/tests.rs | 44 +++++++- apps/desktop-tauri/src-tauri/src/powertoys.rs | 2 + .../src-tauri/src/tray_bridge.rs | 1 + .../src-tauri/src/usage_metric.rs | 1 + .../src/components/MenuCardDetails.tsx | 51 ++++++++- .../providers/sections/UsageSection.test.tsx | 23 ++++ .../providers/sections/UsageSection.tsx | 33 +++++- apps/desktop-tauri/src/types/bridge.ts | 11 ++ rust/src/cli/usage.rs | 100 +++++++++++++++++- rust/src/core/usage_snapshot.rs | 50 +++++++++ 13 files changed, 338 insertions(+), 6 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 8cd5e91924..1a86eab7b4 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -178,6 +178,17 @@ pub struct SubscriptionMetadataSnapshot { pub renews_at: Option, } +/// Display-only provider inventory. Redemption identifiers never cross the +/// bridge. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderInventoryItemSnapshot { + pub id: String, + pub title: String, + pub available_count: u32, + pub next_expires_at: Option, +} + /// A frontend-friendly snapshot of one provider's usage data. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -200,6 +211,8 @@ pub struct ProviderUsageSnapshot { pub tertiary: Option, #[serde(default)] pub extra_rate_windows: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub inventory: Vec, #[serde(default)] pub cost: Option, #[serde(default)] @@ -384,6 +397,16 @@ impl ProviderUsageSnapshot { window: RateWindowSnapshot::from_rate_window(&extra.window), }) .collect(), + inventory: result + .inventory + .iter() + .map(|item| ProviderInventoryItemSnapshot { + id: item.id.clone(), + title: item.title.clone(), + available_count: item.available_count, + next_expires_at: item.next_expires_at.map(|date| date.to_rfc3339()), + }) + .collect(), cost: result.cost.as_ref().map(|c| CostSnapshotBridge { used: c.used, limit: c.limit, @@ -461,6 +484,7 @@ impl ProviderUsageSnapshot { tertiary: None, tertiary_label: None, extra_rate_windows: Vec::new(), + inventory: Vec::new(), cost: None, plan_name: None, account_email: None, diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs index 7690a5b2aa..18b1def39b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs @@ -26,6 +26,7 @@ pub struct ProviderDetail { pub model_specific: Option, pub tertiary: Option, pub extra_rate_windows: Vec, + pub inventory: Vec, // Cost / pace. pub cost: Option, @@ -91,6 +92,7 @@ pub(crate) fn build_provider_detail(provider_id: &str) -> Result::from_timestamp(1_900_000_000, 0).unwrap(); + let result = ProviderFetchResult::new( + codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(12.0)), + "web", + ) + .with_inventory_item(ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count: 2, + next_expires_at: Some(expiry), + }); + let metadata = instantiate_provider(ProviderId::Grok).metadata().clone(); + let snapshot = + ProviderUsageSnapshot::from_fetch_result(ProviderId::Grok, &metadata, &result, None); + + assert_eq!(snapshot.inventory.len(), 1); + assert_eq!(snapshot.inventory[0].available_count, 2); + assert_eq!( + snapshot.inventory[0].next_expires_at.as_deref(), + Some("2030-03-17T17:46:40+00:00") + ); + let serialized = serde_json::to_string(&snapshot).unwrap(); + assert!(serialized.contains("reset-credits")); + assert!(!serialized.contains("coupon-token-secret")); +} + #[test] fn provider_cache_is_fresh_inside_stale_window() { assert!(super::is_provider_cache_fresh( @@ -1057,6 +1085,7 @@ fn provider_cache_upsert_replaces_existing_provider() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(10.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "CLI".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1083,6 +1112,7 @@ fn provider_cache_prunes_disabled_providers() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(10.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "CLI".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1116,6 +1146,7 @@ fn hiding_codex_spark_rows_preserves_other_extra_usage() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(10.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "CLI".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1149,6 +1180,7 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1184,6 +1216,7 @@ fn codex_transient_transport_failure_helper_uses_typed_policy() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1218,6 +1251,7 @@ fn claude_repeated_auth_failure_surfaces_error() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1259,6 +1293,7 @@ fn claude_cloudflare_challenge_retains_prior_usage_while_surfaceing_guidance() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1311,6 +1346,7 @@ fn claude_cloudflare_challenge_keeps_prior_usage_when_guidance_surfaces() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(42.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "Web".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1360,6 +1396,7 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(17.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "CLI".to_string(), has_successful_claude_cli_quota: true, pace_authoritative: true, @@ -1405,6 +1442,7 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() { usage: codexbar::core::UsageSnapshot::new(codexbar::core::RateWindow::new(17.0)), cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1575,6 +1613,7 @@ fn japanese_provider_snapshot_localizes_weekly_label() { usage, cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -1607,6 +1646,7 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() { usage, cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: "OAuth".to_string(), has_successful_claude_cli_quota: false, pace_authoritative: true, diff --git a/apps/desktop-tauri/src-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 4f5ce88465..f38f080161 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -195,6 +195,7 @@ mod tests { tertiary: None, tertiary_label: None, extra_rate_windows: Vec::new(), + inventory: Vec::new(), cost: None, plan_name: Some("Team".to_string()), account_email: Some("dev@example.com".to_string()), @@ -243,6 +244,7 @@ mod tests { tertiary: None, tertiary_label: None, extra_rate_windows: Vec::new(), + inventory: Vec::new(), cost: None, plan_name: None, account_email: None, diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index 00ba7aff36..0ec2e19b85 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -1080,6 +1080,7 @@ mod tests { }), tertiary_label: None, extra_rate_windows: Vec::new(), + inventory: Vec::new(), cost: cost.map(|(used, limit)| crate::commands::CostSnapshotBridge { used, limit: Some(limit), diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 73d642534c..697b94d7a2 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -296,6 +296,7 @@ mod tests { tertiary: None, tertiary_label: None, extra_rate_windows: Vec::new(), + inventory: Vec::new(), cost: None, plan_name: None, account_email: None, diff --git a/apps/desktop-tauri/src/components/MenuCardDetails.tsx b/apps/desktop-tauri/src/components/MenuCardDetails.tsx index 4d0b47d24d..8225a10817 100644 --- a/apps/desktop-tauri/src/components/MenuCardDetails.tsx +++ b/apps/desktop-tauri/src/components/MenuCardDetails.tsx @@ -3,6 +3,7 @@ import type { CostSummaryDisplayStyle, DailyCostPoint, PaceSnapshot, + ProviderInventoryItem, ProviderChartData, ProviderLocalUsageSummary, ProviderUsageSnapshot, @@ -413,6 +414,7 @@ function MetricRow({ export interface MenuCardPresence { hasMetrics: boolean; + hasInventory: boolean; hasCost: boolean; hasPace: boolean; hasCharts: boolean; @@ -456,6 +458,7 @@ export function describeCard( const localUsage = provider.error ? null : chartData?.localUsage ?? null; const wayfinderUsage = isWayfinder ? provider.wayfinderUsage : null; const hasMetrics = visibleMetrics.length > 0; + const hasInventory = !provider.error && (provider.inventory?.length ?? 0) > 0; const hasCost = !!provider.cost && (costSummaryDisplayStyle !== "hidden" || provider.cost.alwaysVisible === true); @@ -465,9 +468,16 @@ export function describeCard( !!provider.pace; const hasDetails = !provider.error && - (hasMetrics || hasCost || hasPace || hasCharts || !!localUsage || !!wayfinderUsage); + (hasMetrics || + hasInventory || + hasCost || + hasPace || + hasCharts || + !!localUsage || + !!wayfinderUsage); return { hasMetrics, + hasInventory, hasCost, hasPace, hasCharts, @@ -505,6 +515,7 @@ export default function MenuCardDetails({ const { hasMetrics, + hasInventory, hasCost, hasPace, hasCharts, @@ -540,6 +551,18 @@ export default function MenuCardDetails({ )} + {!provider.error && hasInventory && ( +
+ {provider.inventory?.map((item) => ( + + ))} +
+ )} + {wayfinderUsage && } {hasMetrics && hasCost &&
} @@ -712,3 +735,29 @@ export default function MenuCardDetails({
); } + +function InventoryItemRow({ + item, + resetTimeRelative, +}: { + item: ProviderInventoryItem; + resetTimeRelative: boolean; +}) { + const formattedExpiry = useFormattedResetTime( + item.nextExpiresAt, + null, + resetTimeRelative, + "expires", + ); + + return ( +
+ {item.title}: {item.availableCount} available + {formattedExpiry && ( + + {formattedExpiry} + + )} +
+ ); +} diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx index 1a4807c882..c32df0cf15 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.test.tsx @@ -102,4 +102,27 @@ describe("UsageSection", () => { expect(label.parentElement).toHaveTextContent("No active 5h session"); expect(label.parentElement?.querySelector(".provider-usage-bar__track")).toBeNull(); }); + + it("renders discrete inventory without turning it into a quota bar", async () => { + const detail = provider(); + detail.session = null; + detail.extraRateWindows = []; + detail.inventory = [ + { + id: "reset-credits", + title: "Limit Reset Credits", + availableCount: 2, + nextExpiresAt: "2099-01-01T00:00:00Z", + }, + ]; + + const { container } = render( + + key} /> + , + ); + + expect(await screen.findByText(/Limit Reset Credits: 2 available/)).toBeInTheDocument(); + expect(container.querySelector(".provider-usage-bar__track")).toBeNull(); + }); }); diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index a51c294ed2..64a3fa7b6a 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -1,4 +1,5 @@ import type { + ProviderInventoryItem, ProviderDetail, RateWindowSnapshot, } from "../../../../types/bridge"; @@ -60,7 +61,8 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { }); } - if (bars.length === 0) { + const inventory = provider.inventory ?? []; + if (bars.length === 0 && inventory.length === 0) { return null; } @@ -76,10 +78,39 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { t={t} /> ))} + {inventory.map((item) => ( + + ))} ); } +function InventoryRow({ + item, + resetTimeRelative, +}: { + item: ProviderInventoryItem; + resetTimeRelative: boolean; +}) { + const formattedExpiry = useFormattedResetTime( + item.nextExpiresAt, + null, + resetTimeRelative, + "expires", + ); + + return ( +
+ {item.title}: {item.availableCount} available + {formattedExpiry && {formattedExpiry}} +
+ ); +} + function UsageBar({ label, rate, diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 865625efeb..963779ae2f 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -592,6 +592,13 @@ export interface SubscriptionMetadataSnapshot { renewsAt: string | null; } +export interface ProviderInventoryItem { + id: string; + title: string; + availableCount: number; + nextExpiresAt: string | null; +} + /** Backend-classified provider availability state (camelCase serde on the bridge). */ export type ProviderStateKind = | "ready" @@ -618,6 +625,8 @@ export interface ProviderUsageSnapshot { title: string; window: RateWindowSnapshot; }>; + /** Display-only discrete provider inventory; never used as quota math. */ + inventory?: ProviderInventoryItem[]; cost: CostSnapshotBridge | null; planName: string | null; accountEmail: string | null; @@ -892,6 +901,8 @@ export interface ProviderDetail { title: string; window: RateWindowSnapshot; }>; + /** Display-only discrete provider inventory; never used as quota math. */ + inventory?: ProviderInventoryItem[]; cost: CostSnapshotBridge | null; pace: PaceSnapshot | null; diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index e0c19618d9..25b988e5c0 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -1,11 +1,13 @@ //! Usage command implementation +use chrono::{DateTime, Utc}; use clap::Args; use serde::Serialize; use crate::core::{ - CostSnapshot, FetchContext, ProviderFetchResult, ProviderId, RateWindow, SourceMode, - TokenAccountStore, TokenAccountSupport, UsagePace, UsageSnapshot, instantiate_provider, + CostSnapshot, FetchContext, ProviderFetchResult, ProviderId, ProviderInventoryItem, RateWindow, + SourceMode, TokenAccountStore, TokenAccountSupport, UsagePace, UsageSnapshot, + instantiate_provider, }; use crate::settings::ApiKeys; use crate::status::{ProviderStatus as StatusInfo, StatusLevel, fetch_provider_status}; @@ -396,6 +398,23 @@ fn render_json_result( }); } + if !result.inventory.is_empty() { + json_result["inventory"] = serde_json::Value::Array( + result + .inventory + .iter() + .map(|item| { + serde_json::json!({ + "id": &item.id, + "title": &item.title, + "availableCount": item.available_count, + "nextExpiresAt": item.next_expires_at.map(|date| date.to_rfc3339()), + }) + }) + .collect(), + ); + } + if let Some(s) = status { json_result["status"] = serde_json::json!({ "level": format!("{:?}", s.level).to_lowercase(), @@ -460,6 +479,7 @@ pub fn render_text_with_status( append_status_line(&mut lines, status); append_account_lines(&mut lines, &result.usage); append_usage_window_lines(&mut lines, &result.usage, &metadata, use_color); + append_inventory_lines(&mut lines, &result.inventory); append_cost_line(&mut lines, result.cost.as_ref()); lines.join("\n") @@ -579,6 +599,38 @@ fn append_usage_window_lines( } } +fn append_inventory_lines(lines: &mut Vec, inventory: &[ProviderInventoryItem]) { + if inventory.is_empty() { + return; + } + let now = Utc::now(); + for item in inventory { + lines.push(format!( + " {}: {} available", + item.title, item.available_count + )); + if let Some(expires_at) = item.next_expires_at { + lines.push(format!( + " Next expires in {}", + format_inventory_countdown(expires_at, now) + )); + } + } +} + +fn format_inventory_countdown(expires_at: DateTime, now: DateTime) -> String { + let seconds = expires_at.signed_duration_since(now).num_seconds(); + if seconds <= 0 { + return "now".to_string(); + } + let minutes = (seconds + 59) / 60; + if minutes >= 24 * 60 { + format!("{}d {}h", minutes / (24 * 60), (minutes / 60) % 24) + } else { + format!("{}h {}m", minutes / 60, minutes % 60) + } +} + fn append_window_line(lines: &mut Vec, label: &str, window: &RateWindow, use_color: bool) { if window.is_informational { let description = window.reset_description.as_deref().unwrap_or("unavailable"); @@ -935,6 +987,50 @@ mod tests { ); } + #[test] + fn inventory_is_rendered_in_full_text_but_not_brief_text() { + let result = fetch_result(UsageSnapshot::new(RateWindow::new(10.0))).with_inventory_item( + ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count: 2, + next_expires_at: Some(Utc::now() + chrono::Duration::hours(3)), + }, + ); + + let full = render_text_with_status(ProviderId::Grok, &result, None, false); + let brief = render_brief_text(ProviderId::Grok, &result); + + assert!(full.contains("Limit Reset Credits: 2 available")); + assert!(full.contains("Next expires in")); + assert!(!brief.contains("Limit Reset Credits")); + } + + #[test] + fn json_inventory_is_additive_and_contains_no_redemption_token() { + let result = fetch_result(UsageSnapshot::new(RateWindow::new(10.0))).with_inventory_item( + ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count: 1, + next_expires_at: None, + }, + ); + + let json = render_json_result(ProviderId::Grok, result, None); + assert_eq!(json["inventory"][0]["availableCount"], 1); + assert!( + serde_json::to_string(&json) + .unwrap() + .contains("reset-credits") + ); + assert!( + !serde_json::to_string(&json) + .unwrap() + .contains("coupon-token-secret") + ); + } + #[test] fn secondary_label_override_is_shared_by_full_and_brief_renderers() { let result = fetch_result( diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 3cfcd2ccf2..b0a1440814 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -85,6 +85,21 @@ pub struct NamedRateWindow { pub usage_known: bool, } +/// One display-only item of provider-issued discrete inventory. +/// +/// This is deliberately separate from [`RateWindow`]: inventory does not +/// represent a percentage quota and must not participate in quota arithmetic, +/// tray metric selection, pace, notifications, or auto-resume decisions. +/// Provider-specific redemption identifiers stay private to the provider +/// parser and never enter this type. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ProviderInventoryItem { + pub id: String, + pub title: String, + pub available_count: u32, + pub next_expires_at: Option>, +} + fn named_rate_window_usage_known_default() -> bool { true } @@ -585,6 +600,13 @@ pub struct ProviderFetchResult { #[serde(skip_serializing_if = "Option::is_none")] pub wayfinder_usage: Option, + /// Transient non-quota inventory for provider-specific display. + /// + /// The field is intentionally skipped by serde: it belongs to the current + /// fetch and must not change persisted `ProviderFetchResult` JSON. + #[serde(skip)] + pub inventory: Vec, + /// Label describing the data source (e.g., "oauth", "web", "cli") pub source_label: String, @@ -613,6 +635,7 @@ impl ProviderFetchResult { usage, cost: None, wayfinder_usage: None, + inventory: Vec::new(), source_label: source_label.into(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -652,6 +675,12 @@ impl ProviderFetchResult { self.wayfinder_usage = Some(usage); self } + + /// Attach one display-only inventory item without exposing redemption IDs. + pub fn with_inventory_item(mut self, item: ProviderInventoryItem) -> Self { + self.inventory.push(item); + self + } } #[cfg(test)] @@ -669,6 +698,27 @@ mod tests { ); } + #[test] + fn fetch_result_inventory_is_transient_and_not_serialized() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let expiry = DateTime::::from_timestamp(1_900_000_000, 0).unwrap(); + let result = + ProviderFetchResult::new(usage, "api").with_inventory_item(ProviderInventoryItem { + id: "reset-credits".to_string(), + title: "Limit Reset Credits".to_string(), + available_count: 2, + next_expires_at: Some(expiry), + }); + + assert_eq!(result.inventory.len(), 1); + let encoded = serde_json::to_value(&result).unwrap(); + assert!(encoded.get("inventory").is_none()); + assert!(encoded.get("reset-credits").is_none()); + + let decoded: ProviderFetchResult = serde_json::from_value(encoded).unwrap(); + assert!(decoded.inventory.is_empty()); + } + #[test] fn cost_snapshot_ignores_non_finite_values() { let cost = CostSnapshot::new(f64::NAN, "USD", "Monthly").with_limit(f64::INFINITY); From 4468703a92a7166e84ad2a44290af2ffe3d0d01b Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sat, 19 Sep 2026 22:47:04 +0700 Subject: [PATCH 02/12] Add transient provider detail carrier --- .../src-tauri/src/commands/bridge.rs | 37 ++++ .../src-tauri/src/commands/provider_detail.rs | 3 + .../src-tauri/src/commands/providers.rs | 1 + .../src-tauri/src/commands/tests.rs | 29 ++- apps/desktop-tauri/src-tauri/src/powertoys.rs | 2 + .../src-tauri/src/tray_bridge.rs | 1 + .../src-tauri/src/usage_metric.rs | 1 + .../src/components/MenuCardDetails.tsx | 37 ++++ .../providers/sections/UsageSection.tsx | 28 ++- apps/desktop-tauri/src/types/bridge.ts | 18 ++ rust/src/cli/usage.rs | 62 ++++++ rust/src/core/usage_snapshot.rs | 179 ++++++++++++++++++ 12 files changed, 394 insertions(+), 4 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 1a86eab7b4..96f05a366a 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -189,6 +189,25 @@ pub struct ProviderInventoryItemSnapshot { pub next_expires_at: Option, } +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderDisplayProgressSnapshot { + pub used: f64, + pub total: f64, +} + +/// Display-only provider detail row. It never participates in quota math or +/// core persistence and contains values validated by the provider carrier. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderDisplayDetailSnapshot { + pub id: String, + pub title: String, + pub value: String, + pub secondary_value: Option, + pub progress: Option, +} + /// A frontend-friendly snapshot of one provider's usage data. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(rename_all = "camelCase")] @@ -213,6 +232,8 @@ pub struct ProviderUsageSnapshot { pub extra_rate_windows: Vec, #[serde(default, skip_serializing_if = "Vec::is_empty")] pub inventory: Vec, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub display_details: Vec, #[serde(default)] pub cost: Option, #[serde(default)] @@ -407,6 +428,21 @@ impl ProviderUsageSnapshot { next_expires_at: item.next_expires_at.map(|date| date.to_rfc3339()), }) .collect(), + display_details: result + .display_details() + .map(|detail| ProviderDisplayDetailSnapshot { + id: detail.id().to_string(), + title: detail.title().to_string(), + value: detail.value().to_string(), + secondary_value: detail.secondary_value().map(ToOwned::to_owned), + progress: detail + .progress() + .map(|progress| ProviderDisplayProgressSnapshot { + used: progress.used(), + total: progress.total(), + }), + }) + .collect(), cost: result.cost.as_ref().map(|c| CostSnapshotBridge { used: c.used, limit: c.limit, @@ -485,6 +521,7 @@ impl ProviderUsageSnapshot { tertiary_label: None, extra_rate_windows: Vec::new(), inventory: Vec::new(), + display_details: Vec::new(), cost: None, plan_name: None, account_email: None, diff --git a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs index 18b1def39b..df5dd7d9c6 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_detail.rs @@ -27,6 +27,7 @@ pub struct ProviderDetail { pub tertiary: Option, pub extra_rate_windows: Vec, pub inventory: Vec, + pub display_details: Vec, // Cost / pace. pub cost: Option, @@ -93,6 +94,7 @@ pub(crate) fn build_provider_detail(provider_id: &str) -> Result 0; const hasInventory = !provider.error && (provider.inventory?.length ?? 0) > 0; + const hasDisplayDetails = !provider.error && (provider.displayDetails?.length ?? 0) > 0; const hasCost = !!provider.cost && (costSummaryDisplayStyle !== "hidden" || provider.cost.alwaysVisible === true); @@ -470,6 +473,7 @@ export function describeCard( !provider.error && (hasMetrics || hasInventory || + hasDisplayDetails || hasCost || hasPace || hasCharts || @@ -478,6 +482,7 @@ export function describeCard( return { hasMetrics, hasInventory, + hasDisplayDetails, hasCost, hasPace, hasCharts, @@ -516,6 +521,7 @@ export default function MenuCardDetails({ const { hasMetrics, hasInventory, + hasDisplayDetails, hasCost, hasPace, hasCharts, @@ -563,6 +569,14 @@ export default function MenuCardDetails({ )} + {!provider.error && hasDisplayDetails && ( +
+ {provider.displayDetails?.map((detail, index) => ( + + ))} +
+ )} + {wayfinderUsage && } {hasMetrics && hasCost &&
} @@ -761,3 +775,26 @@ function InventoryItemRow({
); } + +function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { + const progress = detail.progress; + const progressPercent = progress && Number.isFinite(progress.used) && Number.isFinite(progress.total) && progress.total > 0 + ? Math.max(0, Math.min(100, (progress.used / progress.total) * 100)) + : null; + + return ( +
+
+ {detail.title}: {detail.value} + {detail.secondaryValue && ( + {detail.secondaryValue} + )} +
+ {progressPercent != null && ( +
+
+
+ )} +
+ ); +} diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index 64a3fa7b6a..84323eab68 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -1,4 +1,5 @@ import type { + ProviderDisplayDetail, ProviderInventoryItem, ProviderDetail, RateWindowSnapshot, @@ -62,7 +63,8 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { } const inventory = provider.inventory ?? []; - if (bars.length === 0 && inventory.length === 0) { + const displayDetails = provider.displayDetails ?? []; + if (bars.length === 0 && inventory.length === 0 && displayDetails.length === 0) { return null; } @@ -85,6 +87,9 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { resetTimeRelative={resetTimeRelative} /> ))} + {displayDetails.map((detail, index) => ( + + ))} ); } @@ -111,6 +116,27 @@ function InventoryRow({ ); } +function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { + const progress = detail.progress; + const progressPercent = progress && Number.isFinite(progress.used) && Number.isFinite(progress.total) && progress.total > 0 + ? Math.max(0, Math.min(100, (progress.used / progress.total) * 100)) + : null; + + return ( +
+
+ {detail.title}: {detail.value} + {detail.secondaryValue && {detail.secondaryValue}} +
+ {progressPercent != null && ( +
+
+
+ )} +
+ ); +} + function UsageBar({ label, rate, diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index 963779ae2f..7fd4af94d1 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -599,6 +599,20 @@ export interface ProviderInventoryItem { nextExpiresAt: string | null; } +export interface ProviderDisplayProgress { + used: number; + total: number; +} + +/** Transient provider detail row; it is display-only and never quota math. */ +export interface ProviderDisplayDetail { + id: string; + title: string; + value: string; + secondaryValue: string | null; + progress: ProviderDisplayProgress | null; +} + /** Backend-classified provider availability state (camelCase serde on the bridge). */ export type ProviderStateKind = | "ready" @@ -627,6 +641,8 @@ export interface ProviderUsageSnapshot { }>; /** Display-only discrete provider inventory; never used as quota math. */ inventory?: ProviderInventoryItem[]; + /** Provider-specific display rows; never used as quota math or persistence. */ + displayDetails?: ProviderDisplayDetail[]; cost: CostSnapshotBridge | null; planName: string | null; accountEmail: string | null; @@ -903,6 +919,8 @@ export interface ProviderDetail { }>; /** Display-only discrete provider inventory; never used as quota math. */ inventory?: ProviderInventoryItem[]; + /** Provider-specific display rows; never used as quota math or persistence. */ + displayDetails?: ProviderDisplayDetail[]; cost: CostSnapshotBridge | null; pace: PaceSnapshot | null; diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 25b988e5c0..9697b883c7 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -415,6 +415,28 @@ fn render_json_result( ); } + if result.display_details().next().is_some() { + json_result["details"] = serde_json::Value::Array( + result + .display_details() + .map(|detail| { + serde_json::json!({ + "id": detail.id(), + "title": detail.title(), + "value": detail.value(), + "secondaryValue": detail.secondary_value(), + "progress": detail.progress().map(|progress| { + serde_json::json!({ + "used": progress.used(), + "total": progress.total(), + }) + }), + }) + }) + .collect(), + ); + } + if let Some(s) = status { json_result["status"] = serde_json::json!({ "level": format!("{:?}", s.level).to_lowercase(), @@ -480,6 +502,7 @@ pub fn render_text_with_status( append_account_lines(&mut lines, &result.usage); append_usage_window_lines(&mut lines, &result.usage, &metadata, use_color); append_inventory_lines(&mut lines, &result.inventory); + append_display_detail_lines(&mut lines, result.display_details()); append_cost_line(&mut lines, result.cost.as_ref()); lines.join("\n") @@ -618,6 +641,29 @@ fn append_inventory_lines(lines: &mut Vec, inventory: &[ProviderInventor } } +fn append_display_detail_lines<'a>( + lines: &mut Vec, + details: impl IntoIterator, +) { + for detail in details { + let secondary = detail + .secondary_value() + .map(|value| format!(" ({value})")) + .unwrap_or_default(); + let progress = detail + .progress() + .map(|value| format!(" [{:.2}/{:.2}]", value.used(), value.total())) + .unwrap_or_default(); + lines.push(format!( + " {}: {}{}{}", + detail.title(), + detail.value(), + secondary, + progress + )); + } +} + fn format_inventory_countdown(expires_at: DateTime, now: DateTime) -> String { let seconds = expires_at.signed_duration_since(now).num_seconds(); if seconds <= 0 { @@ -1031,6 +1077,22 @@ mod tests { ); } + #[test] + fn display_details_are_rendered_in_full_text_and_json() { + let result = fetch_result(UsageSnapshot::new(RateWindow::new(10.0))).with_display_detail( + crate::core::ProviderDisplayDetail::new("credits", "Used this cycle", "12") + .with_secondary_value("Monthly refill: 100") + .with_progress(12.0, 100.0), + ); + + let full = render_text_with_status(ProviderId::Grok, &result, None, false); + let json = render_json_result(ProviderId::Grok, result, None); + + assert!(full.contains("Used this cycle: 12 (Monthly refill: 100) [12.00/100.00]")); + assert_eq!(json["details"][0]["title"], "Used this cycle"); + assert_eq!(json["details"][0]["progress"]["total"], 100.0); + } + #[test] fn secondary_label_override_is_shared_by_full_and_brief_renderers() { let result = fetch_result( diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index b0a1440814..86df6c0caa 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -100,6 +100,121 @@ pub struct ProviderInventoryItem { pub next_expires_at: Option>, } +/// One transient provider detail row for display surfaces. +/// +/// These rows are intentionally separate from quota windows and inventory: +/// providers may report credit balances, subscription metadata, or other +/// values that must be shown without becoming quota math or persisted core +/// fetch state. The builder validates compact, display-safe values before a +/// row enters a fetch result; the desktop bridge may then export those rows +/// as part of its current display snapshot. +#[derive(Debug, Clone, PartialEq)] +pub struct ProviderDisplayDetail { + id: String, + title: String, + value: String, + secondary_value: Option, + progress: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq)] +pub struct ProviderDisplayProgress { + used: f64, + total: f64, +} + +impl ProviderDisplayDetail { + pub fn new(id: impl Into, title: impl Into, value: impl Into) -> Self { + Self { + id: id.into(), + title: title.into(), + value: value.into(), + secondary_value: None, + progress: None, + } + } + + pub fn with_secondary_value(mut self, value: impl Into) -> Self { + self.secondary_value = Some(value.into()); + self + } + + pub fn with_progress(mut self, used: f64, total: f64) -> Self { + if used.is_finite() && total.is_finite() && used >= 0.0 && total > 0.0 { + self.progress = Some(ProviderDisplayProgress { used, total }); + } + self + } + + pub fn id(&self) -> &str { + &self.id + } + + pub fn title(&self) -> &str { + &self.title + } + + pub fn value(&self) -> &str { + &self.value + } + + pub fn secondary_value(&self) -> Option<&str> { + self.secondary_value.as_deref() + } + + pub fn progress(&self) -> Option { + self.progress + } + + fn is_display_safe(&self) -> bool { + is_display_safe_text(&self.id, 64) + && is_display_safe_text(&self.title, 128) + && is_display_safe_text(&self.value, 512) + && self + .secondary_value + .as_deref() + .is_none_or(|value| is_display_safe_text(value, 512)) + && self.progress.is_none_or(|progress| { + progress.used.is_finite() + && progress.total.is_finite() + && progress.used >= 0.0 + && progress.total > 0.0 + }) + } +} + +impl ProviderDisplayProgress { + pub fn used(&self) -> f64 { + self.used + } + + pub fn total(&self) -> f64 { + self.total + } +} + +fn is_display_safe_text(value: &str, max_len: usize) -> bool { + if value.is_empty() || value.chars().count() > max_len || value.chars().any(char::is_control) { + return false; + } + + let lower = value.to_ascii_lowercase(); + [ + "authorization:", + "bearer ", + "cookie:", + "set-cookie:", + "access_token", + "api_key", + "api-key", + "client_secret", + "refresh_token", + "x-api-key", + ] + .iter() + .all(|marker| !lower.contains(marker)) +} + fn named_rate_window_usage_known_default() -> bool { true } @@ -607,6 +722,12 @@ pub struct ProviderFetchResult { #[serde(skip)] pub inventory: Vec, + /// Transient provider-specific detail rows for display only. They are not + /// serialized by the core result; use [`Self::display_details`] for an + /// explicit surface projection. + #[serde(skip)] + pub display_details: Vec, + /// Label describing the data source (e.g., "oauth", "web", "cli") pub source_label: String, @@ -636,6 +757,7 @@ impl ProviderFetchResult { cost: None, wayfinder_usage: None, inventory: Vec::new(), + display_details: Vec::new(), source_label: source_label.into(), has_successful_claude_cli_quota: false, pace_authoritative: true, @@ -681,6 +803,20 @@ impl ProviderFetchResult { self.inventory.push(item); self } + + /// Attach one transient provider-specific detail row without persisting it. + pub fn with_display_detail(mut self, detail: ProviderDisplayDetail) -> Self { + if detail.is_display_safe() && !self.display_details.iter().any(|row| row.id == detail.id) { + self.display_details.push(detail); + } + self + } + + pub fn display_details(&self) -> impl Iterator { + self.display_details + .iter() + .filter(|detail| detail.is_display_safe()) + } } #[cfg(test)] @@ -719,6 +855,49 @@ mod tests { assert!(decoded.inventory.is_empty()); } + #[test] + fn fetch_result_display_details_are_transient_and_validate_progress() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let result = ProviderFetchResult::new(usage, "web").with_display_detail( + ProviderDisplayDetail::new("credits", "Used this cycle", "12") + .with_secondary_value("Monthly refill: 100") + .with_progress(12.0, 100.0), + ); + + let details: Vec<_> = result.display_details().collect(); + assert_eq!(details.len(), 1); + assert!(details[0].progress().is_some()); + assert!( + ProviderDisplayDetail::new("invalid", "Invalid", "value") + .with_progress(f64::NAN, 1.0) + .progress + .is_none() + ); + let encoded = serde_json::to_value(&result).unwrap(); + assert!(encoded.get("display_details").is_none()); + } + + #[test] + fn display_details_reject_secret_markers_and_duplicate_ids() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let result = ProviderFetchResult::new(usage, "web") + .with_display_detail(ProviderDisplayDetail::new("credits", "Credits", "12")) + .with_display_detail(ProviderDisplayDetail::new( + "credits", + "Credits duplicate", + "13", + )) + .with_display_detail(ProviderDisplayDetail::new( + "secret", + "Authorization", + "Bearer hidden", + )); + + let details: Vec<_> = result.display_details().collect(); + assert_eq!(details.len(), 1); + assert_eq!(details[0].value(), "12"); + } + #[test] fn cost_snapshot_ignores_non_finite_values() { let cost = CostSnapshot::new(f64::NAN, "USD", "Monthly").with_limit(f64::INFINITY); From 2a8a4c5dd0c8f0876c1dbd15d28169e16f701220 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sun, 20 Sep 2026 02:18:32 +0700 Subject: [PATCH 03/12] Port Replicate billing provider --- CHANGELOG.md | 1 + README.md | 1 + .../src/commands/provider_settings.rs | 17 + .../src-tauri/src/commands/tests.rs | 21 + .../icons/ProviderIcon-replicate.svg | 1 + .../src/components/providers/providerIcons.ts | 3 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + docs/PROVIDERS.md | 10 + rust/src/cli/serve/dashboard/icons.rs | 4 + .../icons/ProviderIcon-replicate.svg | 1 + rust/src/core/provider.rs | 10 +- rust/src/core/provider_factory.rs | 7 +- rust/src/core/token_accounts.rs | 8 + rust/src/providers/mod.rs | 2 + rust/src/providers/replicate/mod.rs | 808 ++++++++++++++++++ 15 files changed, 891 insertions(+), 4 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-replicate.svg create mode 100644 rust/src/cli/serve/dashboard/icons/ProviderIcon-replicate.svg create mode 100644 rust/src/providers/replicate/mod.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a1d2ed3f74..4309b941a1 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. --- 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/provider_settings.rs b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs index 829018b850..65a2cb6b89 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/provider_settings.rs @@ -185,6 +185,7 @@ fn cookie_source_provider(provider_id: &str) -> Option ProviderId::Sakana, "notion" => ProviderId::Notion, "grok" => ProviderId::Grok, + "replicate" => ProviderId::Replicate, _ => return None, }) } @@ -690,6 +691,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/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index a81fd4e41f..846aa1018b 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(); @@ -1704,6 +1718,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 0b554eba9b..c0696258ae 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -50,6 +50,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"; @@ -133,6 +134,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), @@ -207,6 +209,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/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index 3281d8c070..c3078fddad 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -67,5 +67,6 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["sub2api", "sub2api"], ["qwencloud", "Qwen Cloud"], ["notion", "Notion AI"], + ["replicate", "Replicate"], ["meta", "Meta"], ]; 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/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 36fc7c3326..c40c6f0ffc 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -249,6 +249,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/core/provider.rs b/rust/src/core/provider.rs index 55ab70d391..6883f01bd8 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -84,6 +84,7 @@ pub enum ProviderId { Fireworks, #[serde(alias = "metaspark")] Meta, + Replicate, } impl ProviderId { @@ -161,6 +162,7 @@ impl ProviderId { ProviderId::Xai, ProviderId::Fireworks, ProviderId::Meta, + ProviderId::Replicate, ] } @@ -238,6 +240,7 @@ impl ProviderId { ProviderId::QwenCloud => "qwen-cloud", ProviderId::Notion => "notion", ProviderId::Xai => "xai", + ProviderId::Replicate => "replicate", } } @@ -317,6 +320,7 @@ impl ProviderId { ProviderId::QwenCloud => "Qwen Cloud", ProviderId::Notion => "Notion AI", ProviderId::Xai => "xAI", + ProviderId::Replicate => "Replicate", } } @@ -356,6 +360,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, @@ -492,6 +497,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, } } @@ -976,6 +982,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Xai => "#8E8E93", ProviderId::Fireworks => "#F25B1C", ProviderId::Meta => "#0467DF", + ProviderId::Replicate => "#000000", } } @@ -990,7 +997,7 @@ mod tests { #[test] fn test_provider_id_all() { let all = ProviderId::all(); - assert_eq!(all.len(), 71); + assert_eq!(all.len(), 72); assert!(all.contains(&ProviderId::Claude)); assert!(all.contains(&ProviderId::Codex)); assert!(all.contains(&ProviderId::Fireworks)); @@ -1042,6 +1049,7 @@ mod tests { assert!(all.contains(&ProviderId::Notion)); assert!(all.contains(&ProviderId::Xai)); assert!(all.contains(&ProviderId::Meta)); + assert!(all.contains(&ProviderId::Replicate)); } #[test] diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 093b5cf089..d02f04e821 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -18,9 +18,9 @@ use crate::providers::{ MiMoProvider, MiniMaxProvider, MistralProvider, 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, + ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, + VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, + XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -96,6 +96,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 a8ae172b17..eae78570aa 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/providers/mod.rs b/rust/src/providers/mod.rs index 1e32aaecd5..988e0c25be 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -63,6 +63,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; @@ -136,6 +137,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; diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs new file mode 100644 index 0000000000..1513827061 --- /dev/null +++ b/rust/src/providers/replicate/mod.rs @@ -0,0 +1,808 @@ +//! 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, reuses the shared browser-cookie cache, or imports the +//! `replicate.com` browser session. It never uses a Replicate API token as a +//! website credential and never logs cookie material. + +use async_trait::async_trait; +use chrono::{DateTime, 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::browser::cookie_cache::{CookieHeaderCache, CookieHeaderEntry}; +use crate::core::{ + CostSnapshot, FetchContext, 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, + }, + 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 mut observed = CookieHeaderCache::load(ProviderId::Replicate); + if let Some(cached) = observed.as_ref() { + match self + .fetch_with_cookie(&cached.cookie_header, &cached.source_label) + .await + { + Ok(result) => return Ok(result), + Err(error) if is_authentication_failure(&error) => { + clear_cache_if_current(cached); + observed = None; + } + Err(error) => return Err(error), + } + } + + let imported = crate::providers::browser_cookie_header(&["replicate.com"])?; + let normalized = normalize_cookie_header(&imported).ok_or(ProviderError::NoCookies)?; + let result = self.fetch_with_cookie(&normalized, "browser").await?; + store_cache_if_current(observed.as_ref(), &normalized, "browser"); + Ok(result) + } + + async fn fetch_auto(&self, ctx: &FetchContext) -> Result { + if let Some(cookie_header) = ctx.manual_cookie_header.as_deref() { + return self.fetch_with_cookie(cookie_header, "manual").await; + } + 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 => self.fetch_auto(ctx).await, + SourceMode::Web => { + if let Some(cookie_header) = ctx.manual_cookie_header.as_deref() { + self.fetch_with_cookie(cookie_header, "manual").await + } else { + self.fetch_browser_cookie().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 + } +} + +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 scripts = 0usize; + let lower = body.to_ascii_lowercase(); + let mut cursor = 0usize; + while cursor < lower.len() && scripts < MAX_REACT_NODES { + 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 + " bool { + let lower = attributes.to_ascii_lowercase(); + let name = name.to_ascii_lowercase(); + let expected = expected.to_ascii_lowercase(); + let Some(mut cursor) = lower.find(&name) else { + return false; + }; + while cursor < lower.len() { + let before = cursor + .checked_sub(1) + .and_then(|index| lower.as_bytes().get(index)); + let after = lower.as_bytes().get(cursor + name.len()); + if before.is_none_or(|value| !value.is_ascii_alphanumeric()) + && after.is_none_or(|value| !value.is_ascii_alphanumeric()) + { + let rest = lower[cursor + name.len()..].trim_start(); + if let Some(rest) = rest.strip_prefix('=') { + let rest = rest.trim_start(); + if let Some(rest) = rest.strip_prefix('"') { + return rest + .split_once('"') + .is_some_and(|(value, _)| value == expected); + } + if let Some(rest) = rest.strip_prefix('\'') { + return rest + .split_once('\'') + .is_some_and(|(value, _)| value == expected); + } + } + } + let next = cursor + name.len(); + let Some(relative) = lower[next..].find(&name) else { + break; + }; + cursor = next + relative; + } + false +} + +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() => { + DateTime::parse_from_rfc3339(value) + .map(|end| end.with_timezone(&Utc) > now) + .unwrap_or(false) + } + _ => 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_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 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:.0}s." + ))); + } + if status == StatusCode::REQUEST_TIMEOUT || status.is_server_error() { + return Err(ProviderError::Other(format!( + "Replicate billing is unavailable; retry after {retry_after:.0}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) +} + +fn same_cache_entry(left: &CookieHeaderEntry, right: &CookieHeaderEntry) -> bool { + left.cookie_header == right.cookie_header + && left.source_label == right.source_label + && left.stored_at == right.stored_at +} + +fn clear_cache_if_current(expected: &CookieHeaderEntry) { + if CookieHeaderCache::load(ProviderId::Replicate) + .as_ref() + .is_some_and(|current| same_cache_entry(current, expected)) + { + CookieHeaderCache::clear(ProviderId::Replicate); + } +} + +fn store_cache_if_current( + expected: Option<&CookieHeaderEntry>, + cookie_header: &str, + source_label: &str, +) { + let current = CookieHeaderCache::load(ProviderId::Replicate); + let unchanged = match (expected, current.as_ref()) { + (None, None) => true, + (Some(expected), Some(current)) => same_cache_entry(expected, current), + _ => false, + }; + if unchanged + && let Err(error) = + CookieHeaderCache::store(ProviderId::Replicate, cookie_header, source_label) + { + tracing::debug!(%error, "Replicate: failed to persist browser cookie cache"); + } +} + +#[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_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: Vec<_> = result.display_details().collect(); + 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 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") + ); + } + + #[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()); + } +} From e78e247bc02df926123579cee816ca4c35081921 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sun, 20 Sep 2026 02:41:42 +0700 Subject: [PATCH 04/12] Harden Replicate browser authentication --- .../src-tauri/src/commands/providers.rs | 17 +- .../src-tauri/src/commands/tests.rs | 15 ++ apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 2 +- rust/src/browser/cookies.rs | 54 +++++++ rust/src/providers/mod.rs | 13 ++ rust/src/providers/replicate/mod.rs | 153 +++++++++++------- 6 files changed, 194 insertions(+), 60 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index d7798d8cab..752e407ba5 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -121,7 +121,13 @@ pub(crate) fn build_fetch_context( "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) + let source_mode = if id == ProviderId::Replicate && cookie_header.is_none() { + // Replicate's manual mode means an explicitly pasted + // sessionid cookie. Preserve an empty sentinel so the + // provider fails closed instead of the generic web + // fallback importing a different browser account. + SourceMode::Web + } else if (has_kimi_code_api_key || has_opencodego_api_key) && usage_source == SourceMode::Auto { SourceMode::Auto @@ -137,7 +143,14 @@ pub(crate) fn build_fetch_context( } else { SourceMode::Cli }; - (source_mode, cookie_header) + ( + source_mode, + if id == ProviderId::Replicate && cookie_header.is_none() { + Some(String::new()) + } else { + cookie_header + }, + ) } // `browser` is accepted as a legacy alias from older settings. "auto" | "browser" | "web" => { diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 846aa1018b..179238ea27 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -469,6 +469,21 @@ 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_eq!(ctx.manual_cookie_header.as_deref(), Some("")); +} + #[test] fn fetch_context_codex_manual_cookie_never_forces_unsupported_web() { // Default cookie source is "manual". Pasting a chatgpt.com cookie used to flip diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 05c4156ce0..063d552af3 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -28,7 +28,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/rust/src/browser/cookies.rs b/rust/src/browser/cookies.rs index 0d21c4251d..6c940f8403 100755 --- a/rust/src/browser/cookies.rs +++ b/rust/src/browser/cookies.rs @@ -717,6 +717,60 @@ pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> Err(CookieError::NotFound(domain.to_string())) } +/// Get cookie-header candidates from every detected browser that has readable +/// cookies for a domain. +/// +/// The older `get_cookie_header` helper intentionally stops at the first +/// browser with any matching cookie. 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> { + use super::detection::BrowserDetector; + + let browsers = BrowserDetector::detect_all(); + if browsers.is_empty() { + return Err(CookieError::BrowserNotInstalled); + } + + let mut candidates = Vec::new(); + let mut abe_error_seen = false; + + for browser in browsers { + match CookieExtractor::extract_for_domain(&browser, domain) { + Ok(cookies) => { + let header = CookieExtractor::build_cookie_header(&cookies); + if !header.trim().is_empty() { + candidates.push((browser.browser_type, header)); + } + } + Err(CookieError::AppBoundEncryption) => { + abe_error_seen = true; + tracing::debug!( + browser = %browser.browser_type.display_name(), + "App-Bound Encryption prevented cookie candidate extraction" + ); + } + Err(error) => { + tracing::debug!( + browser = %browser.browser_type.display_name(), + %error, + "Failed to extract browser cookie candidates" + ); + } + } + } + + if candidates.is_empty() && abe_error_seen { + return Err(CookieError::AppBoundEncryption); + } + if candidates.is_empty() { + return Err(CookieError::NotFound(domain.to_string())); + } + Ok(candidates) +} + /// 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/providers/mod.rs b/rust/src/providers/mod.rs index 988e0c25be..236dafca50 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -160,6 +160,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) +} + pub(crate) fn browser_cookies_for_domain( domain: &str, ) -> Result, crate::core::ProviderError> { diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs index 1513827061..96cf96383f 100644 --- a/rust/src/providers/replicate/mod.rs +++ b/rust/src/providers/replicate/mod.rs @@ -3,12 +3,12 @@ //! 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, reuses the shared browser-cookie cache, or imports the -//! `replicate.com` browser session. It never uses a Replicate API token as a -//! website credential and never logs cookie material. +//! 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, Utc}; +use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc}; use futures::StreamExt; use reqwest::{Client, StatusCode, Url, header::HeaderMap}; use serde_json::Value; @@ -16,7 +16,6 @@ use std::collections::VecDeque; use std::time::Duration; use tokio::time::timeout; -use crate::browser::cookie_cache::{CookieHeaderCache, CookieHeaderEntry}; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, @@ -162,26 +161,25 @@ impl ReplicateProvider { } async fn fetch_browser_cookie(&self) -> Result { - let mut observed = CookieHeaderCache::load(ProviderId::Replicate); - if let Some(cached) = observed.as_ref() { - match self - .fetch_with_cookie(&cached.cookie_header, &cached.source_label) - .await - { + 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) => { - clear_cache_if_current(cached); - observed = None; + authentication_failed = true; } Err(error) => return Err(error), } } - let imported = crate::providers::browser_cookie_header(&["replicate.com"])?; - let normalized = normalize_cookie_header(&imported).ok_or(ProviderError::NoCookies)?; - let result = self.fetch_with_cookie(&normalized, "browser").await?; - store_cache_if_current(observed.as_ref(), &normalized, "browser"); - Ok(result) + if authentication_failed { + Err(ProviderError::AuthRequired) + } else { + Err(ProviderError::NoCookies) + } } async fn fetch_auto(&self, ctx: &FetchContext) -> Result { @@ -433,9 +431,7 @@ fn parse_current_invoice(body: &str, now: DateTime) -> Result true, Some(Value::String(value)) if !value.trim().is_empty() => { - DateTime::parse_from_rfc3339(value) - .map(|end| end.with_timezone(&Utc) > now) - .unwrap_or(false) + parse_invoice_end(value).is_some_and(|end| end > now) } _ => false, } @@ -448,6 +444,31 @@ fn parse_current_invoice(body: &str, now: DateTime) -> Result 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() @@ -547,6 +568,15 @@ fn normalize_cookie_header(raw: &str) -> Option { }) } +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); @@ -558,12 +588,12 @@ fn validate_status(status: StatusCode, headers: &HeaderMap) -> Result<(), Provid ); if status == StatusCode::TOO_MANY_REQUESTS { return Err(ProviderError::Other(format!( - "Replicate rate limit reached; retry after {retry_after:.0}s." + "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:.0}s." + "Replicate billing is unavailable; retry after {retry_after:.3}s." ))); } if !status.is_success() { @@ -592,40 +622,6 @@ fn is_authentication_failure(error: &ProviderError) -> bool { matches!(error, ProviderError::AuthRequired) } -fn same_cache_entry(left: &CookieHeaderEntry, right: &CookieHeaderEntry) -> bool { - left.cookie_header == right.cookie_header - && left.source_label == right.source_label - && left.stored_at == right.stored_at -} - -fn clear_cache_if_current(expected: &CookieHeaderEntry) { - if CookieHeaderCache::load(ProviderId::Replicate) - .as_ref() - .is_some_and(|current| same_cache_entry(current, expected)) - { - CookieHeaderCache::clear(ProviderId::Replicate); - } -} - -fn store_cache_if_current( - expected: Option<&CookieHeaderEntry>, - cookie_header: &str, - source_label: &str, -) { - let current = CookieHeaderCache::load(ProviderId::Replicate); - let unchanged = match (expected, current.as_ref()) { - (None, None) => true, - (Some(expected), Some(current)) => same_cache_entry(expected, current), - _ => false, - }; - if unchanged - && let Err(error) = - CookieHeaderCache::store(ProviderId::Replicate, cookie_header, source_label) - { - tracing::debug!(%error, "Replicate: failed to persist browser cookie cache"); - } -} - #[cfg(test)] mod tests { use super::*; @@ -711,6 +707,27 @@ mod tests { ); } + #[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(); @@ -775,6 +792,19 @@ mod tests { 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(); @@ -798,6 +828,15 @@ mod tests { .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] From 760c2227897a3c2d3f48519980e6ebe086b114d2 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sun, 20 Sep 2026 23:13:11 +0700 Subject: [PATCH 05/12] Refactor Replicate provider policy into provider-siloed hooks --- .../src-tauri/src/commands/mod.rs | 6 +- .../src-tauri/src/commands/providers.rs | 166 +++++++++--------- .../src-tauri/src/commands/tests.rs | 3 +- rust/src/browser/cookies.rs | 113 ++++++------ rust/src/cli/diagnose.rs | 1 + rust/src/cli/guard.rs | 1 + rust/src/cli/hooks.rs | 1 + rust/src/cli/serve/dashboard/source.rs | 2 + rust/src/cli/serve/data.rs | 1 + rust/src/cli/usage.rs | 1 + rust/src/core/provider.rs | 26 +++ rust/src/providers/replicate/mod.rs | 132 ++++++++------ 12 files changed, 252 insertions(+), 201 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/mod.rs b/apps/desktop-tauri/src-tauri/src/commands/mod.rs index 12678ed50c..f9c34e95b4 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/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 752e407ba5..47ed877d31 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -84,94 +84,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 id == ProviderId::Replicate && cookie_header.is_none() { - // Replicate's manual mode means an explicitly pasted - // sessionid cookie. Preserve an empty sentinel so the - // provider fails closed instead of the generic web - // fallback importing a different browser account. - SourceMode::Web - } else 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, - if id == ProviderId::Replicate && cookie_header.is_none() { - Some(String::new()) - } else { - 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 @@ -219,6 +216,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), api_region: (!api_region.is_empty()).then_some(api_region), diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 179238ea27..71b9c5d045 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -481,7 +481,8 @@ fn fetch_context_replicate_empty_manual_fails_closed_without_browser_import() { ); assert_eq!(ctx.source_mode, SourceMode::Web); - assert_eq!(ctx.manual_cookie_header.as_deref(), Some("")); + assert!(ctx.manual_cookie_header.is_none()); + assert!(ctx.manual_cookie_missing); } #[test] diff --git a/rust/src/browser/cookies.rs b/rust/src/browser/cookies.rs index 6c940f8403..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,21 +699,44 @@ 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); } @@ -720,55 +746,28 @@ pub fn get_cookies_for_domain(domain: &str) -> Result, CookieError> /// Get cookie-header candidates from every detected browser that has readable /// cookies for a domain. /// -/// The older `get_cookie_header` helper intentionally stops at the first -/// browser with any matching cookie. 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. +/// 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> { - use super::detection::BrowserDetector; + let scan = extract_domain_candidates(domain)?; - let browsers = BrowserDetector::detect_all(); - if browsers.is_empty() { - return Err(CookieError::BrowserNotInstalled); - } - - let mut candidates = Vec::new(); - let mut abe_error_seen = false; - - for browser in browsers { - match CookieExtractor::extract_for_domain(&browser, domain) { - Ok(cookies) => { - let header = CookieExtractor::build_cookie_header(&cookies); - if !header.trim().is_empty() { - candidates.push((browser.browser_type, header)); - } - } - Err(CookieError::AppBoundEncryption) => { - abe_error_seen = true; - tracing::debug!( - browser = %browser.browser_type.display_name(), - "App-Bound Encryption prevented cookie candidate extraction" - ); - } - Err(error) => { - tracing::debug!( - browser = %browser.browser_type.display_name(), - %error, - "Failed to extract browser cookie candidates" - ); - } - } - } + let headers = scan + .candidates + .into_iter() + .map(|(browser, cookies)| (browser, CookieExtractor::build_cookie_header(&cookies))) + .filter(|(_, header)| !header.trim().is_empty()) + .collect::>(); - if candidates.is_empty() && abe_error_seen { + if headers.is_empty() && scan.abe_error_seen { return Err(CookieError::AppBoundEncryption); } - if candidates.is_empty() { + if headers.is_empty() { return Err(CookieError::NotFound(domain.to_string())); } - Ok(candidates) + Ok(headers) } /// Get a cookie header string for a domain diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index f97ccda8b7..a24edd07ef 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -171,6 +171,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 cd3cb5693e..e0b69c93aa 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, api_region: None, diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index bfec1e12ff..9186348871 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -289,6 +289,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()), api_region: (!region.is_empty()).then(|| region.to_string()), diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 61b491c16d..149a711705 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -155,6 +155,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, api_region: None, @@ -274,6 +275,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, api_region: None, diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 9697b883c7..49325336c8 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -236,6 +236,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, api_region: None, diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 6883f01bd8..5dcb202788 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -689,6 +689,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, @@ -721,6 +726,7 @@ impl Default for FetchContext { web_timeout: 60, verbose: false, manual_cookie_header: None, + manual_cookie_missing: false, api_key: None, workspace_id: None, api_region: None, @@ -740,6 +746,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 { @@ -782,6 +798,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 diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs index 96cf96383f..ce4d45b8c6 100644 --- a/rust/src/providers/replicate/mod.rs +++ b/rust/src/providers/replicate/mod.rs @@ -17,8 +17,9 @@ use std::time::Duration; use tokio::time::timeout; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderDisplayDetail, ProviderError, - ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + CostSnapshot, FetchContext, ManualEmptyCookiePolicy, Provider, ProviderDisplayDetail, + ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, RateWindow, SourceMode, + UsageSnapshot, }; const BILLING_URL: &str = "https://replicate.com/account/billing"; @@ -182,10 +183,25 @@ impl ReplicateProvider { } } - async fn fetch_auto(&self, ctx: &FetchContext) -> Result { + /// 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 } } @@ -208,14 +224,7 @@ impl Provider for ReplicateProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { - SourceMode::Auto => self.fetch_auto(ctx).await, - SourceMode::Web => { - if let Some(cookie_header) = ctx.manual_cookie_header.as_deref() { - self.fetch_with_cookie(cookie_header, "manual").await - } else { - self.fetch_browser_cookie().await - } - } + SourceMode::Auto | SourceMode::Web => self.fetch_with_cookie_source(ctx).await, source => Err(ProviderError::UnsupportedSource(source)), } } @@ -231,6 +240,10 @@ impl Provider for ReplicateProvider { 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> { @@ -278,10 +291,14 @@ fn account_endpoint(account: &ReplicateAccount, suffix: &str) -> Result Result { - let mut scripts = 0usize; + let mut scanned = 0usize; let lower = body.to_ascii_lowercase(); let mut cursor = 0usize; - while cursor < lower.len() && scripts < MAX_REACT_NODES { + while cursor < lower.len() { + scanned += 1; + if scanned > MAX_REACT_NODES { + break; + } let Some(relative_start) = lower[cursor..].find(" Result break; }; let close = content_start + relative_close; - let attributes = &body[after_name..tag_end]; - if has_script_attribute(attributes, "id", "react-component-props") - && has_script_attribute(attributes, "type", "application/json") + if parse_script_attributes(&body[after_name..tag_end]).is_some_and(|attrs| { + attrs + .iter() + .any(|(name, value)| name == "id" && value == "react-component-props") + && attrs + .iter() + .any(|(name, value)| name == "type" && value == "application/json") + }) && let Ok(value) = serde_json::from_str::(&body[content_start..close]) + && let Some(account) = find_account_value(&value) { - scripts += 1; - if let Ok(value) = serde_json::from_str::(&body[content_start..close]) - && let Some(account) = find_account_value(&value) - { - return Ok(account); - } + return Ok(account); } cursor = close + " Result )) } -fn has_script_attribute(attributes: &str, name: &str, expected: &str) -> bool { - let lower = attributes.to_ascii_lowercase(); - let name = name.to_ascii_lowercase(); - let expected = expected.to_ascii_lowercase(); - let Some(mut cursor) = lower.find(&name) else { - return false; - }; - while cursor < lower.len() { - let before = cursor - .checked_sub(1) - .and_then(|index| lower.as_bytes().get(index)); - let after = lower.as_bytes().get(cursor + name.len()); - if before.is_none_or(|value| !value.is_ascii_alphanumeric()) - && after.is_none_or(|value| !value.is_ascii_alphanumeric()) - { - let rest = lower[cursor + name.len()..].trim_start(); - if let Some(rest) = rest.strip_prefix('=') { - let rest = rest.trim_start(); - if let Some(rest) = rest.strip_prefix('"') { - return rest - .split_once('"') - .is_some_and(|(value, _)| value == expected); - } - if let Some(rest) = rest.strip_prefix('\'') { - return rest - .split_once('\'') - .is_some_and(|(value, _)| value == expected); - } +/// Parse `name="value"` / `name='value'` pairs from a raw script tag attribute +/// string. Unquoted and malformed attributes are skipped, matching the lenient +/// reading the previous hand-rolled matcher accepted for the target tags. +fn parse_script_attributes(raw: &str) -> 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; } - let next = cursor + name.len(); - let Some(relative) = lower[next..].find(&name) else { - break; - }; - cursor = next + relative; } - false + Some(attrs) } fn find_account_value(root: &Value) -> Option { From 1fd80ba0134498919669b5cd000ef48d2c759e98 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:22:53 +0700 Subject: [PATCH 06/12] Fix Replicate merge integration with main --- .../src-tauri/src/commands/bridge.rs | 10 ---------- rust/src/core/usage_snapshot.rs | 15 --------------- rust/src/providers/replicate/mod.rs | 1 + 3 files changed, 1 insertion(+), 25 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index ae2aa0daa5..62d1976398 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -433,16 +433,6 @@ impl ProviderUsageSnapshot { next_expires_at: item.next_expires_at.map(|date| date.to_rfc3339()), }) .collect(), - inventory: result - .inventory - .iter() - .map(|item| ProviderInventoryItemSnapshot { - id: item.id.clone(), - title: item.title.clone(), - available_count: item.available_count, - next_expires_at: item.next_expires_at.map(|date| date.to_rfc3339()), - }) - .collect(), display_details: result .display_details() .map(|detail| ProviderDisplayDetailSnapshot { diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 4e3ec9c023..5f6cf0f8b2 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -105,21 +105,6 @@ pub struct ProviderInventoryItem { pub next_expires_at: Option>, } -/// One display-only item of provider-issued discrete inventory. -/// -/// This is deliberately separate from [`RateWindow`]: inventory does not -/// represent a percentage quota and must not participate in quota arithmetic, -/// tray metric selection, pace, notifications, or auto-resume decisions. -/// Provider-specific redemption identifiers stay private to the provider -/// parser and never enter this type. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct ProviderInventoryItem { - pub id: String, - pub title: String, - pub available_count: u32, - pub next_expires_at: Option>, -} - /// One transient provider detail row for display surfaces. /// /// These rows are intentionally separate from quota windows and inventory: diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs index ce4d45b8c6..285e53dbce 100644 --- a/rust/src/providers/replicate/mod.rs +++ b/rust/src/providers/replicate/mod.rs @@ -74,6 +74,7 @@ impl ReplicateProvider { 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) From 69e8ba4184477a5cc98321a49184e378f8bfdf68 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 03:27:24 +0700 Subject: [PATCH 07/12] Fix duplicate closing brace in provider inventory test --- apps/desktop-tauri/src-tauri/src/commands/tests.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index 54fe4369fe..a8457a3e90 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -10,8 +10,6 @@ use crate::surface_target::SurfaceTarget; use codexbar::core::{ FetchContext, ProviderAccountData, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, ProviderInventoryItem, SourceMode, TokenAccount, instantiate_provider, - FetchContext, ProviderAccountData, ProviderError, ProviderFetchResult, ProviderId, - ProviderInventoryItem, SourceMode, TokenAccount, instantiate_provider, }; use codexbar::host::session::launch_block_reason; use codexbar::settings::{ApiKeys, Language, ManualCookies, Settings}; @@ -1031,7 +1029,6 @@ fn provider_inventory_maps_to_the_bridge_without_token_ids() { .with_secondary_value("Monthly refill: 100") .with_progress(12.0, 100.0), ); - }); let metadata = instantiate_provider(ProviderId::Grok).metadata().clone(); let snapshot = ProviderUsageSnapshot::from_fetch_result(ProviderId::Grok, &metadata, &result, None); From 7ac422c221d21c2b27369736b3bb13d92172f6d7 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 07:40:53 +0700 Subject: [PATCH 08/12] Deduplicate display detail type after refresh merge --- rust/src/core/usage_snapshot.rs | 113 +------------------------------- 1 file changed, 2 insertions(+), 111 deletions(-) diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 16e4d3c8f5..9e3fbdc662 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -4,8 +4,7 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use super::RateWindow; -use crate::core::ProviderDisplayDetail; - +use super::{ProviderDisplayDetail, ProviderDisplayProgress}; /// Subscription dates explicitly reported by an authenticated provider /// dashboard or subscription endpoint. /// @@ -106,99 +105,6 @@ pub struct ProviderInventoryItem { pub next_expires_at: Option>, } -/// One transient provider detail row for display surfaces. -/// -/// These rows are intentionally separate from quota windows and inventory: -/// providers may report credit balances, subscription metadata, or other -/// values that must be shown without becoming quota math or persisted core -/// fetch state. The builder validates compact, display-safe values before a -/// row enters a fetch result; the desktop bridge may then export those rows -/// as part of its current display snapshot. -#[derive(Debug, Clone, PartialEq)] -pub struct ProviderDisplayDetail { - id: String, - title: String, - value: String, - secondary_value: Option, - progress: Option, -} - -#[derive(Debug, Clone, Copy, PartialEq)] -pub struct ProviderDisplayProgress { - used: f64, - total: f64, -} - -impl ProviderDisplayDetail { - pub fn new(id: impl Into, title: impl Into, value: impl Into) -> Self { - Self { - id: id.into(), - title: title.into(), - value: value.into(), - secondary_value: None, - progress: None, - } - } - - pub fn with_secondary_value(mut self, value: impl Into) -> Self { - self.secondary_value = Some(value.into()); - self - } - - pub fn with_progress(mut self, used: f64, total: f64) -> Self { - if used.is_finite() && total.is_finite() && used >= 0.0 && total > 0.0 { - self.progress = Some(ProviderDisplayProgress { used, total }); - } - self - } - - pub fn id(&self) -> &str { - &self.id - } - - pub fn title(&self) -> &str { - &self.title - } - - pub fn value(&self) -> &str { - &self.value - } - - pub fn secondary_value(&self) -> Option<&str> { - self.secondary_value.as_deref() - } - - pub fn progress(&self) -> Option { - self.progress - } - - fn is_display_safe(&self) -> bool { - is_display_safe_text(&self.id, 64) - && is_display_safe_text(&self.title, 128) - && is_display_safe_text(&self.value, 512) - && self - .secondary_value - .as_deref() - .is_none_or(|value| is_display_safe_text(value, 512)) - && self.progress.is_none_or(|progress| { - progress.used.is_finite() - && progress.total.is_finite() - && progress.used >= 0.0 - && progress.total > 0.0 - }) - } -} - -impl ProviderDisplayProgress { - pub fn used(&self) -> f64 { - self.used - } - - pub fn total(&self) -> f64 { - self.total - } -} - fn is_display_safe_text(value: &str, max_len: usize) -> bool { if value.is_empty() || value.chars().count() > max_len || value.chars().any(char::is_control) { return false; @@ -818,22 +724,7 @@ impl ProviderFetchResult { pub fn with_inventory_item(mut self, item: ProviderInventoryItem) -> Self { self.inventory.push(item); self - } - - /// Attach one transient provider-specific detail row without persisting it. - pub fn with_display_detail(mut self, detail: ProviderDisplayDetail) -> Self { - if detail.is_display_safe() && !self.display_details.iter().any(|row| row.id == detail.id) { - self.display_details.push(detail); - } - self - } - - pub fn display_details(&self) -> impl Iterator { - self.display_details - .iter() - .filter(|detail| detail.is_display_safe()) - } -} + }} #[cfg(test)] mod tests { From 6d3ec3d13513a9a85027a6818d9a6835c8ab0df2 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 08:00:01 +0700 Subject: [PATCH 09/12] Fix Muse brand color union resolution --- rust/src/core/provider.rs | 2 +- rust/src/core/provider_factory.rs | 8 ++++---- rust/src/core/usage_snapshot.rs | 27 +++------------------------ rust/src/providers/replicate/mod.rs | 2 +- 4 files changed, 9 insertions(+), 30 deletions(-) diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 1bf8a2f57e..2f69a7ac6d 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -1049,7 +1049,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Xai => "#8E8E93", ProviderId::Fireworks => "#F25B1C", ProviderId::Meta => "#0467DF", - ProviderId::Muse => "#0688E1", + ProviderId::Muse => "#0668E1", ProviderId::Replicate => "#000000", } } diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 800f5cf1a6..a88bc9b4ec 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -17,10 +17,10 @@ use crate::providers::{ KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, LongCatProvider, ManusProvider, MetaProvider, MiMoProvider, MiniMaxProvider, MistralProvider, MuseProvider, NanoGPTProvider, NeuralwattProvider, NotionProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, - OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PoeProvider, ReplicateProvider, QoderProvider, - QwenCloudProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, - VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, - XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, + OpenCodeProvider, OpenRouterProvider, 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`]. diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 9e3fbdc662..8c7c883340 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -3,8 +3,8 @@ use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; +use super::ProviderDisplayDetail; use super::RateWindow; -use super::{ProviderDisplayDetail, ProviderDisplayProgress}; /// Subscription dates explicitly reported by an authenticated provider /// dashboard or subscription endpoint. /// @@ -105,28 +105,6 @@ pub struct ProviderInventoryItem { pub next_expires_at: Option>, } -fn is_display_safe_text(value: &str, max_len: usize) -> bool { - if value.is_empty() || value.chars().count() > max_len || value.chars().any(char::is_control) { - return false; - } - - let lower = value.to_ascii_lowercase(); - [ - "authorization:", - "bearer ", - "cookie:", - "set-cookie:", - "access_token", - "api_key", - "api-key", - "client_secret", - "refresh_token", - "x-api-key", - ] - .iter() - .all(|marker| !lower.contains(marker)) -} - fn named_rate_window_usage_known_default() -> bool { true } @@ -724,7 +702,8 @@ impl ProviderFetchResult { pub fn with_inventory_item(mut self, item: ProviderInventoryItem) -> Self { self.inventory.push(item); self - }} + } +} #[cfg(test)] mod tests { diff --git a/rust/src/providers/replicate/mod.rs b/rust/src/providers/replicate/mod.rs index 285e53dbce..1bcebe3aa0 100644 --- a/rust/src/providers/replicate/mod.rs +++ b/rust/src/providers/replicate/mod.rs @@ -796,7 +796,7 @@ mod tests { 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: Vec<_> = result.display_details().collect(); + 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"); From 2cd730be73f62772091e42fe4ed3d83be4aae14c Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:20:04 +0700 Subject: [PATCH 10/12] Resolve leftover conflict marker in UsageSection imports --- .../src/surfaces/settings/providers/sections/UsageSection.tsx | 3 --- 1 file changed, 3 deletions(-) diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index 8109175fe1..5b7a97b0e7 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -1,9 +1,6 @@ import type { ProviderDisplayDetail, -<<<<<<< HEAD -======= ProviderInventoryItem, ->>>>>>> origin/main ProviderDetail, RateWindowSnapshot, } from "../../../../types/bridge"; From 1dc5519748ff1f11a9801ba3148e3788174af3df Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 11:20:46 +0700 Subject: [PATCH 11/12] Adopt main's UsageSection, dropping committed conflict markers --- .../providers/sections/UsageSection.tsx | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx index 5b7a97b0e7..d5686dd69c 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/UsageSection.tsx @@ -94,10 +94,6 @@ export function UsageSection({ provider, resetTimeRelative, t }: Props) { lineClassName="provider-usage-inventory" /> ))} -<<<<<<< HEAD - {displayDetails.map((detail, index) => ( - -======= {displayDetails.map((detail) => ( ->>>>>>> origin/main ))} ); } -function DisplayDetailRow({ detail }: { detail: ProviderDisplayDetail }) { - const progress = detail.progress; - const progressPercent = progress && Number.isFinite(progress.used) && Number.isFinite(progress.total) && progress.total > 0 - ? Math.max(0, Math.min(100, (progress.used / progress.total) * 100)) - : null; - - return ( -
-
- {detail.title}: {detail.value} - {detail.secondaryValue && {detail.secondaryValue}} -
- {progressPercent != null && ( -
-
-
- )} -
- ); -} - function UsageBar({ label, rate, From ffa5e546c713ade75e93e80f1c6736b46540e97d Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:13:08 +0700 Subject: [PATCH 12/12] Set provider count and format --- rust/src/core/provider.rs | 2 +- rust/src/core/provider_factory.rs | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 718b0e878f..af71e750c9 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -1079,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)); diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index b6a75cb813..ca62d1a85a 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -20,8 +20,8 @@ use crate::providers::{ OpenAIApiProvider, OpenCodeGoProvider, OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PoeProvider, QoderProvider, QwenCloudProvider, ReplicateProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, VeniceProvider, - VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, - ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, + VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, XaiProvider, ZaiProvider, + ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`].