From 73d3e1cb1b81e2110a05d0f9724653719804e103 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sun, 20 Sep 2026 01:01:51 +0700 Subject: [PATCH 1/2] Port Muse Code subscription usage --- .../src-tauri/src/commands/tests.rs | 27 + .../providers/icons/ProviderIcon-muse.svg | 1 + .../src/components/providers/providerIcons.ts | 3 + apps/desktop-tauri/src/surfaces/TrayPanel.tsx | 2 +- .../providers/sections/usageSourcePolicy.ts | 6 + .../desktop-tauri/src/test/providerCatalog.ts | 1 + rust/src/cli/serve/dashboard/icons.rs | 4 + .../dashboard/icons/ProviderIcon-muse.svg | 1 + rust/src/core/provider.rs | 25 +- rust/src/core/provider_factory.rs | 13 +- rust/src/core/token_accounts.rs | 3 +- rust/src/providers/mod.rs | 2 + rust/src/providers/muse/mod.rs | 611 ++++++++++++++++++ 13 files changed, 690 insertions(+), 9 deletions(-) create mode 100644 apps/desktop-tauri/src/components/providers/icons/ProviderIcon-muse.svg create mode 100644 rust/src/cli/serve/dashboard/icons/ProviderIcon-muse.svg create mode 100644 rust/src/providers/muse/mod.rs diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index a81fd4e41f..467bd8f59f 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -304,6 +304,33 @@ fn fetch_context_defaults_to_manual_cookies_without_browser_import() { assert_eq!(ctx.source_mode, SourceMode::Web); } +#[test] +fn fetch_context_muse_uses_oauth_login_lane() { + let settings = Settings::default(); + let cookies = ManualCookies::default(); + let api_keys = ApiKeys::default(); + let token_accounts = HashMap::new(); + + let ctx = super::build_fetch_context( + ProviderId::Muse, + &settings, + &cookies, + &api_keys, + &token_accounts, + ); + + assert_eq!(ctx.source_mode, SourceMode::Auto); + let provider = instantiate_provider(ProviderId::Muse); + assert_eq!( + provider.available_sources(), + vec![SourceMode::Auto, SourceMode::OAuth] + ); + assert!(provider.supports_oauth()); + assert!(!provider.supports_web()); + assert!(!provider.supports_cli()); + assert_eq!(provider.metadata().display_name, "Muse Code"); +} + #[test] fn fetch_context_cursor_cookie_off_stays_cli() { let mut settings = Settings::default(); diff --git a/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-muse.svg b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-muse.svg new file mode 100644 index 0000000000..b1d1810486 --- /dev/null +++ b/apps/desktop-tauri/src/components/providers/icons/ProviderIcon-muse.svg @@ -0,0 +1 @@ + \ 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..f4890d63ed 100644 --- a/apps/desktop-tauri/src/components/providers/providerIcons.ts +++ b/apps/desktop-tauri/src/components/providers/providerIcons.ts @@ -42,6 +42,7 @@ import meta from "./icons/ProviderIcon-meta.svg?raw"; import mimo from "./icons/ProviderIcon-mimo.svg?raw"; import minimax from "./icons/ProviderIcon-minimax.svg?raw"; import mistral from "./icons/ProviderIcon-mistral.svg?raw"; +import muse from "./icons/ProviderIcon-muse.svg?raw"; import notion from "./icons/ProviderIcon-notion.svg?raw"; import xai from "./icons/ProviderIcon-xai.svg?raw"; import ollama from "./icons/ProviderIcon-ollama.svg?raw"; @@ -127,6 +128,7 @@ const RAW: Record = { notion: tint(notion), xai: tint(xai), mistral: tint(mistral), + muse: tint(muse), ollama: tint(ollama), opencode: tint(opencode), opencodego: tint(opencodego), @@ -184,6 +186,7 @@ export const PROVIDER_ICON_REGISTRY: Record = { llmproxy: { id: "llmproxy", brandColor: "#4f46e5", fallbackLetter: "L", svgPath: RAW.llmproxy }, minimax: { id: "minimax", brandColor: "#fe603c", fallbackLetter: "M", svgPath: RAW.minimax }, mistral: { id: "mistral", brandColor: "#ff500f", fallbackLetter: "M", svgPath: RAW.mistral }, + muse: { id: "muse", brandColor: "#0668e1", fallbackLetter: "M", svgPath: RAW.muse }, ollama: { id: "ollama", brandColor: "#8b95b0", fallbackLetter: "○", svgPath: RAW.ollama }, azureopenai: { id: "azureopenai", brandColor: "#0078d4", fallbackLetter: "A" }, t3chat: { id: "t3chat", brandColor: "#8b5cf6", fallbackLetter: "T", svgPath: RAW.t3chat }, diff --git a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx index 05c4156ce0..e1c44098b7 100644 --- a/apps/desktop-tauri/src/surfaces/TrayPanel.tsx +++ b/apps/desktop-tauri/src/surfaces/TrayPanel.tsx @@ -32,7 +32,7 @@ const HAS_DASHBOARD = new Set([ "mimo", "minimax", "mistral", "nanogpt", "notion", "ollama", "openaiapi", "opencode", "opencodego", "openrouter", "perplexity", "qoder", "codebuddy", "sakana", "stepfun", "t3chat", "venice", "vertexai", "warp", "windsurf", - "xai", "zai", "fireworks", "meta", + "xai", "zai", "fireworks", "meta", "muse", ]); /** Provider IDs that have a status page URL in the backend */ const HAS_STATUS_PAGE = new Set([ diff --git a/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts b/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts index 90b88c4ff2..fb53b7d203 100644 --- a/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts +++ b/apps/desktop-tauri/src/surfaces/settings/providers/sections/usageSourcePolicy.ts @@ -40,6 +40,12 @@ const POLICIES: Readonly> = { }, ], }, + muse: { + options: [ + { value: "auto", label: "Auto", description: "Uses the local Muse Code device login." }, + { value: "oauth", label: "Muse Code login", description: "Uses the local Muse Code device-code login only." }, + ], + }, }; export function usageSourcePolicy(providerId: string): UsageSourcePolicy | null { diff --git a/apps/desktop-tauri/src/test/providerCatalog.ts b/apps/desktop-tauri/src/test/providerCatalog.ts index 3281d8c070..3186dfa938 100644 --- a/apps/desktop-tauri/src/test/providerCatalog.ts +++ b/apps/desktop-tauri/src/test/providerCatalog.ts @@ -68,4 +68,5 @@ export const TEST_PROVIDER_CATALOG: Array<[string, string]> = [ ["qwencloud", "Qwen Cloud"], ["notion", "Notion AI"], ["meta", "Meta"], + ["muse", "Muse Code"], ]; diff --git a/rust/src/cli/serve/dashboard/icons.rs b/rust/src/cli/serve/dashboard/icons.rs index 36fc7c3326..3b1274ffef 100644 --- a/rust/src/cli/serve/dashboard/icons.rs +++ b/rust/src/cli/serve/dashboard/icons.rs @@ -209,6 +209,10 @@ static ICONS: &[(&str, &[u8])] = &[ "ProviderIcon-mistral", include_bytes!("icons/ProviderIcon-mistral.svg"), ), + ( + "ProviderIcon-muse", + include_bytes!("icons/ProviderIcon-muse.svg"), + ), ( "ProviderIcon-neuralwatt", include_bytes!("icons/ProviderIcon-neuralwatt.svg"), diff --git a/rust/src/cli/serve/dashboard/icons/ProviderIcon-muse.svg b/rust/src/cli/serve/dashboard/icons/ProviderIcon-muse.svg new file mode 100644 index 0000000000..b1d1810486 --- /dev/null +++ b/rust/src/cli/serve/dashboard/icons/ProviderIcon-muse.svg @@ -0,0 +1 @@ + \ No newline at end of file diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 55ab70d391..45e940bcd2 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, + Muse, } impl ProviderId { @@ -161,6 +162,7 @@ impl ProviderId { ProviderId::Xai, ProviderId::Fireworks, ProviderId::Meta, + ProviderId::Muse, ] } @@ -204,6 +206,7 @@ impl ProviderId { ProviderId::DeepInfra => "deepinfra", ProviderId::Fireworks => "fireworks", ProviderId::Meta => "meta", + ProviderId::Muse => "muse", ProviderId::AiAnd => "aiand", ProviderId::Windsurf => "windsurf", ProviderId::Manus => "manus", @@ -282,6 +285,7 @@ impl ProviderId { ProviderId::DeepInfra => "DeepInfra", ProviderId::Fireworks => "Fireworks", ProviderId::Meta => "Meta", + ProviderId::Muse => "Muse Code", ProviderId::AiAnd => "ai&", ProviderId::Windsurf => "Windsurf", ProviderId::Manus => "Manus", @@ -373,6 +377,7 @@ impl ProviderId { ProviderId::DeepInfra => None, ProviderId::Fireworks => None, ProviderId::Meta => None, + ProviderId::Muse => None, ProviderId::AiAnd => None, ProviderId::Windsurf => None, ProviderId::Doubao => None, @@ -446,6 +451,7 @@ impl ProviderId { "deepseek" | "deep-seek" | "ds" => Some(ProviderId::DeepSeek), "deepinfra" | "deep-infra" | "di" => Some(ProviderId::DeepInfra), "fireworks" | "fireworks-ai" | "fw" => Some(ProviderId::Fireworks), + "muse" | "muse-code" | "muse code" => Some(ProviderId::Muse), "meta" | "metaspark" | "meta-spark" | "muse-spark" | "musespark" | "muse spark" | "meta muse spark" => Some(ProviderId::Meta), "aiand" | "ai&" | "ai-and" | "ai and" => Some(ProviderId::AiAnd), @@ -842,6 +848,8 @@ pub fn cli_name_map() -> HashMap<&'static str, ProviderId> { map.insert("di", ProviderId::DeepInfra); map.insert("fireworks-ai", ProviderId::Fireworks); map.insert("fw", ProviderId::Fireworks); + map.insert("muse-code", ProviderId::Muse); + map.insert("muse code", ProviderId::Muse); map.insert("metaspark", ProviderId::Meta); map.insert("meta-spark", ProviderId::Meta); map.insert("muse-spark", ProviderId::Meta); @@ -976,6 +984,7 @@ pub fn brand_color(id: ProviderId) -> &'static str { ProviderId::Xai => "#8E8E93", ProviderId::Fireworks => "#F25B1C", ProviderId::Meta => "#0467DF", + ProviderId::Muse => "#0668E1", } } @@ -990,7 +999,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 +1051,7 @@ mod tests { assert!(all.contains(&ProviderId::Notion)); assert!(all.contains(&ProviderId::Xai)); assert!(all.contains(&ProviderId::Meta)); + assert!(all.contains(&ProviderId::Muse)); } #[test] @@ -1292,6 +1302,19 @@ mod tests { assert_eq!(brand_color(ProviderId::Meta), "#0467DF"); } + #[test] + fn test_provider_id_muse() { + assert_eq!(ProviderId::Muse.cli_name(), "muse"); + assert_eq!(ProviderId::Muse.display_name(), "Muse Code"); + assert_eq!(ProviderId::Muse.cookie_domain(), None); + assert_eq!(ProviderId::from_cli_name("muse"), Some(ProviderId::Muse)); + assert_eq!( + ProviderId::from_cli_name("muse-code"), + Some(ProviderId::Muse) + ); + assert_eq!(brand_color(ProviderId::Muse), "#0668E1"); + } + #[test] fn test_provider_id_xai() { assert_eq!(ProviderId::Xai.cli_name(), "xai"); diff --git a/rust/src/core/provider_factory.rs b/rust/src/core/provider_factory.rs index 093b5cf089..116e536048 100644 --- a/rust/src/core/provider_factory.rs +++ b/rust/src/core/provider_factory.rs @@ -15,12 +15,12 @@ use crate::providers::{ ElevenLabsProvider, FactoryProvider, FireworksProvider, GeminiProvider, GrokProvider, GroqProvider, InfiniProvider, JetBrainsProvider, KiloProvider, KimiK2Provider, KimiProvider, KiroProvider, LLMProxyProvider, LiteLLMProvider, LongCatProvider, ManusProvider, MetaProvider, - 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, + MiMoProvider, MiniMaxProvider, MistralProvider, MuseProvider, NanoGPTProvider, + NeuralwattProvider, NotionProvider, OllamaProvider, OpenAIApiProvider, OpenCodeGoProvider, + OpenCodeProvider, OpenRouterProvider, PerplexityProvider, PoeProvider, QoderProvider, + QwenCloudProvider, SakanaProvider, StepFunProvider, Sub2ApiProvider, T3ChatProvider, + VeniceProvider, VertexAIProvider, WarpProvider, WayfinderProvider, WindsurfProvider, + XaiProvider, ZaiProvider, ZedProvider, ZenMuxProvider, ZoomMateProvider, }; /// Instantiate the concrete [`Provider`] implementation for a given [`ProviderId`]. @@ -100,6 +100,7 @@ pub fn instantiate(id: ProviderId) -> Box { ProviderId::Xai => Box::new(XaiProvider::new()), ProviderId::Fireworks => Box::new(FireworksProvider::new()), ProviderId::Meta => Box::new(MetaProvider::new()), + ProviderId::Muse => Box::new(MuseProvider::new()), } } diff --git a/rust/src/core/token_accounts.rs b/rust/src/core/token_accounts.rs index a8ae172b17..694a9d65ef 100755 --- a/rust/src/core/token_accounts.rs +++ b/rust/src/core/token_accounts.rs @@ -352,7 +352,8 @@ impl TokenAccountSupport { | ProviderId::Wayfinder | ProviderId::QwenCloud | ProviderId::Fireworks - | ProviderId::Meta => None, + | ProviderId::Meta + | ProviderId::Muse => None, } } diff --git a/rust/src/providers/mod.rs b/rust/src/providers/mod.rs index 1e32aaecd5..dba9aff4cd 100755 --- a/rust/src/providers/mod.rs +++ b/rust/src/providers/mod.rs @@ -50,6 +50,7 @@ pub mod meta; pub mod mimo; pub mod minimax; pub mod mistral; +pub mod muse; pub mod nanogpt; pub mod neuralwatt; pub mod notion; @@ -124,6 +125,7 @@ pub use meta::MetaProvider; pub use mimo::MiMoProvider; pub use minimax::{MiniMaxProvider, MiniMaxRegion}; pub use mistral::MistralProvider; +pub use muse::MuseProvider; pub use nanogpt::NanoGPTProvider; pub use neuralwatt::NeuralwattProvider; pub use notion::NotionProvider; diff --git a/rust/src/providers/muse/mod.rs b/rust/src/providers/muse/mod.rs new file mode 100644 index 0000000000..83bdabbb53 --- /dev/null +++ b/rust/src/providers/muse/mod.rs @@ -0,0 +1,611 @@ +//! Muse Code subscription provider. +//! +//! Muse usage is minted from the Muse CLI's device-code OAuth token (`dca:`). +//! The Windows port reads the CLI auth file without writing or refreshing it; +//! the optional environment override is intended for controlled deployments, +//! not for inference keys (`LLM_` / `LLM|`). + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use reqwest::{Client, StatusCode}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; +use std::path::{Path, PathBuf}; +use tokio::time::{Duration, timeout}; + +use crate::core::{ + FetchContext, Provider, ProviderDisplayDetail, ProviderError, ProviderFetchResult, ProviderId, + ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, +}; + +const MUSE_USAGE_URL: &str = "https://api.meta.ai/muse-code/key"; +const MUSE_DEVICE_TOKEN_ENV: &str = "MUSE_DEVICE_TOKEN"; +const MUSE_AUTH_PATH_ENV: &str = "MUSE_AUTH_PATH"; +const DEVICE_TOKEN_PREFIX: &str = "dca:"; +const REQUEST_TIMEOUT: Duration = Duration::from_secs(15); +const MAX_RESPONSE_BYTES: usize = 512 * 1024; +const WEEKLY_WINDOW_MINUTES: u32 = 7 * 24 * 60; +const MAX_RESET_SECONDS: f64 = 64_092_211_200.0; + +#[derive(Debug, Deserialize)] +struct MuseAuthFile { + providers: Option, +} + +#[derive(Debug, Deserialize)] +struct MuseProviders { + meta: Option, +} + +#[derive(Debug, Deserialize)] +struct MuseMetaCredentials { + mechanism: Option, + #[serde(alias = "accessToken")] + access_token: Option, +} + +pub struct MuseProvider { + metadata: ProviderMetadata, + client: Client, +} + +impl MuseProvider { + pub fn new() -> Self { + Self { + metadata: ProviderMetadata { + id: ProviderId::Muse, + display_name: "Muse Code", + session_label: "5 hours", + weekly_label: "Weekly", + supports_opus: false, + supports_credits: false, + default_enabled: false, + is_primary: false, + dashboard_url: Some("https://dev.meta.ai"), + status_page_url: None, + }, + client: crate::core::credentialed_http_client_builder() + .timeout(REQUEST_TIMEOUT) + .build() + .unwrap_or_else(|_| Client::new()), + } + } + + async fn fetch_api(&self) -> Result { + let token = resolve_device_token()?; + let response = timeout( + REQUEST_TIMEOUT, + self.client + .post(MUSE_USAGE_URL) + .bearer_auth(token) + .header("x-api-version", "1.0.0") + .header("User-Agent", "CodexBar") + .header("Accept", "application/json") + .json(&serde_json::json!({})) + .send(), + ) + .await + .map_err(|_| ProviderError::Timeout)??; + + let status = response.status(); + if status != StatusCode::OK { + return Err(status_error(status)); + } + if response + .content_length() + .is_some_and(|length| length > MAX_RESPONSE_BYTES as u64) + { + return Err(ProviderError::Parse( + "Muse Code returned an oversized response.".to_string(), + )); + } + + let body = response.bytes().await.map_err(ProviderError::Network)?; + if body.len() > MAX_RESPONSE_BYTES { + return Err(ProviderError::Parse( + "Muse Code returned an oversized response.".to_string(), + )); + } + parse_response(&body) + } +} + +impl Default for MuseProvider { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl Provider for MuseProvider { + fn id(&self) -> ProviderId { + ProviderId::Muse + } + + fn metadata(&self) -> &ProviderMetadata { + &self.metadata + } + + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { + match ctx.source_mode { + SourceMode::Auto | SourceMode::OAuth => self.fetch_api().await, + SourceMode::Web | SourceMode::Cli => { + Err(ProviderError::UnsupportedSource(ctx.source_mode)) + } + } + } + + fn available_sources(&self) -> Vec { + vec![SourceMode::Auto, SourceMode::OAuth] + } + + fn supports_oauth(&self) -> bool { + true + } +} + +fn resolve_device_token() -> Result { + let environment: HashMap = std::env::vars().collect(); + let home = dirs::home_dir().ok_or_else(missing_credentials)?; + resolve_device_token_from(&environment, &home) +} + +fn resolve_device_token_from( + environment: &HashMap, + home_directory: &Path, +) -> Result { + if let Some(raw) = environment.get(MUSE_DEVICE_TOKEN_ENV) { + return require_device_token(raw); + } + + let path = auth_file_path(environment, home_directory); + let contents = std::fs::read_to_string(&path).map_err(|_| missing_credentials())?; + let file: MuseAuthFile = serde_json::from_str(&contents).map_err(|_| missing_credentials())?; + let Some(meta) = file.providers.and_then(|providers| providers.meta) else { + return Err(missing_credentials()); + }; + + if let Some(raw) = meta.access_token { + // An invalid inline credential must fail closed. Do not fall through + // to another store and silently switch the selected Muse account. + return require_device_token(&raw); + } + + let mechanism = meta + .mechanism + .as_deref() + .unwrap_or_default() + .trim() + .to_ascii_lowercase(); + if mechanism == "oauth" { + return Err(ProviderError::AuthRequired); + } + Err(missing_credentials()) +} + +fn auth_file_path(environment: &HashMap, home_directory: &Path) -> PathBuf { + environment + .get(MUSE_AUTH_PATH_ENV) + .map(|value| value.trim()) + .filter(|path| !path.is_empty()) + .map(PathBuf::from) + .unwrap_or_else(|| { + home_directory + .join(".config") + .join("muse") + .join("auth.json") + }) +} + +fn require_device_token(raw: &str) -> Result { + let token = raw.trim(); + if token.is_empty() || !token.starts_with(DEVICE_TOKEN_PREFIX) { + return Err(ProviderError::AuthRequired); + } + Ok(token.to_string()) +} + +fn missing_credentials() -> ProviderError { + ProviderError::NotInstalled( + "Muse Code login not found. Run `muse login`, then refresh CodexBar.".to_string(), + ) +} + +fn status_error(status: StatusCode) -> ProviderError { + match status { + StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => ProviderError::AuthRequired, + StatusCode::TOO_MANY_REQUESTS => { + ProviderError::Other("Muse Code usage requests are rate limited.".to_string()) + } + status if status.is_server_error() => { + ProviderError::Other("Muse Code API is unavailable.".to_string()) + } + status => ProviderError::Other(format!("Muse Code API returned HTTP {}.", status.as_u16())), + } +} + +fn parse_response(body: &[u8]) -> Result { + let decoded: Value = serde_json::from_slice(body).map_err(|_| { + ProviderError::Parse( + "Could not parse Muse Code subscription usage: expected JSON".to_string(), + ) + })?; + let root = object(&decoded, "expected a response object")?; + + if optional_bool(root.get("require_payment"), "require_payment")? == Some(true) { + return Err(ProviderError::Other( + "Muse Code requires a payment method. Finish billing at https://dev.meta.ai." + .to_string(), + )); + } + if optional_bool(root.get("is_subs_active"), "is_subs_active")? != Some(true) { + return Err(ProviderError::Other( + "No Muse Code subscription is active on this login.".to_string(), + )); + } + + let usage = object( + root.get("subs_usage") + .ok_or_else(|| parse_failure("missing subs_usage"))?, + "subs_usage", + )?; + let window = object( + usage + .get("window") + .ok_or_else(|| parse_failure("missing subscription window"))?, + "window", + )?; + let weekly = object( + usage + .get("weekly") + .ok_or_else(|| parse_failure("missing weekly window"))?, + "weekly", + )?; + + let duration = positive_safe_minutes(number( + window + .get("window_duration_mins") + .ok_or_else(|| parse_failure("window_duration_mins"))?, + "window_duration_mins", + )?)?; + let primary_percent = clamped_percent(number( + window + .get("used_percent") + .ok_or_else(|| parse_failure("used_percent"))?, + "used_percent", + )?); + let weekly_percent = clamped_percent(number( + weekly + .get("used_percent") + .ok_or_else(|| parse_failure("weekly.used_percent"))?, + "weekly.used_percent", + )?); + let primary_reset = parse_reset(window.get("resets_at"))?; + let weekly_reset = parse_reset(weekly.get("resets_at"))?; + let plan = optional_text(root.get("subs_tier_name"), "subs_tier_name")?; + let email = optional_text(root.get("user_email"), "user_email")?; + + let primary = RateWindow::with_details(primary_percent, Some(duration), primary_reset, None); + let secondary = RateWindow::with_details( + weekly_percent, + Some(WEEKLY_WINDOW_MINUTES), + weekly_reset, + None, + ); + let login_method = plan.clone().unwrap_or_else(|| "Muse login".to_string()); + let mut usage = UsageSnapshot::new(primary) + .with_secondary(secondary) + .with_login_method(login_method); + if let Some(email) = email { + usage = usage.with_email(email); + } + + let mut result = ProviderFetchResult::new(usage, "oauth") + .with_display_detail(ProviderDisplayDetail::new( + "five-hour", + "5 hours", + format_percent(primary_percent), + )) + .with_display_detail(ProviderDisplayDetail::new( + "weekly", + "Weekly", + format_percent(weekly_percent), + )); + if let Some(plan) = plan { + result = result.with_display_detail(ProviderDisplayDetail::new("plan", "Plan", plan)); + } + Ok(result) +} + +fn object<'a>( + value: &'a Value, + field: &str, +) -> Result<&'a serde_json::Map, ProviderError> { + value + .as_object() + .ok_or_else(|| parse_failure(format!("{field} must be an object"))) +} + +fn optional_bool(value: Option<&Value>, field: &str) -> Result, ProviderError> { + match value { + None | Some(Value::Null) => Ok(None), + Some(value) => value + .as_bool() + .map(Some) + .ok_or_else(|| parse_failure(field)), + } +} + +fn number(value: &Value, field: &str) -> Result { + let value = value.as_f64().ok_or_else(|| parse_failure(field))?; + value + .is_finite() + .then_some(value) + .ok_or_else(|| parse_failure(field)) +} + +fn optional_text(value: Option<&Value>, field: &str) -> Result, ProviderError> { + match value { + None | Some(Value::Null) => Ok(None), + Some(value) => value + .as_str() + .map(|text| text.trim().to_string()) + .filter(|text| !text.is_empty()) + .map(Some) + .ok_or_else(|| parse_failure(field)), + } +} + +fn positive_safe_minutes(value: f64) -> Result { + if value <= 0.0 || value.fract() != 0.0 || value > u32::MAX as f64 { + return Err(parse_failure("window_duration_mins")); + } + #[allow( + clippy::cast_possible_truncation, + reason = "value is bounded to u32 range" + )] + let value = value as u32; + Ok(value) +} + +fn clamped_percent(value: f64) -> f64 { + value.clamp(0.0, 100.0) +} + +fn parse_reset(value: Option<&Value>) -> Result>, ProviderError> { + let Some(value) = value else { + return Ok(None); + }; + if value.is_null() { + return Ok(None); + } + let seconds = number(value, "resets_at")?; + if seconds <= 0.0 || seconds > MAX_RESET_SECONDS { + return Ok(None); + } + #[allow( + clippy::cast_possible_truncation, + reason = "reset range is bounded to i64" + )] + let whole_seconds = seconds.trunc() as i64; + Ok(DateTime::::from_timestamp(whole_seconds, 0)) +} + +fn format_percent(percent: f64) -> String { + format!("{percent:.0}%") +} + +fn parse_failure(field: impl Into) -> ProviderError { + ProviderError::Parse(format!( + "Could not parse Muse Code subscription usage: {}", + field.into() + )) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use tempfile::tempdir; + + fn environment(entries: &[(&str, &str)]) -> HashMap { + entries + .iter() + .map(|(key, value)| ((*key).to_string(), (*value).to_string())) + .collect() + } + + fn success_payload() -> Value { + serde_json::json!({ + "require_payment": false, + "is_subs_active": true, + "subs_tier_name": "Pro", + "user_email": "muse@example.com", + "api_key": "LLM_should_not_surface", + "payment_method": "card_should_not_surface", + "subs_usage": { + "window": { + "window_duration_mins": 300, + "used_percent": 96, + "resets_at": 1788580000 + }, + "weekly": { + "used_percent": 40, + "resets_at": 1789000000 + } + } + }) + } + + #[test] + fn provider_metadata_and_sources_are_oauth_only() { + let provider = MuseProvider::new(); + assert_eq!(provider.id(), ProviderId::Muse); + assert_eq!(provider.metadata().display_name, "Muse Code"); + assert!(!provider.metadata().default_enabled); + assert_eq!( + provider.available_sources(), + vec![SourceMode::Auto, SourceMode::OAuth] + ); + assert!(provider.supports_oauth()); + assert!(!provider.supports_web()); + assert!(!provider.supports_cli()); + } + + #[test] + fn auth_file_uses_override_and_requires_device_token_prefix() { + let dir = tempdir().unwrap(); + let path = dir.path().join("custom-auth.json"); + fs::write( + &path, + r#"{"providers":{"meta":{"mechanism":"oauth","access_token":" dca:device-token "}}}"#, + ) + .unwrap(); + let path = path.to_string_lossy().into_owned(); + let env = environment(&[(MUSE_AUTH_PATH_ENV, &path)]); + assert_eq!( + resolve_device_token_from(&env, Path::new("C:\\unused")).unwrap(), + "dca:device-token" + ); + + fs::write( + &path, + r#"{"providers":{"meta":{"mechanism":"oauth","access_token":"LLM-inference-key"}}}"#, + ) + .unwrap(); + assert!(matches!( + resolve_device_token_from(&env, Path::new("C:\\unused")), + Err(ProviderError::AuthRequired) + )); + } + + #[test] + fn auth_file_default_path_and_env_override_are_deterministic() { + let dir = tempdir().unwrap(); + let auth_path = dir.path().join(".config/muse/auth.json"); + fs::create_dir_all(auth_path.parent().unwrap()).unwrap(); + fs::write( + &auth_path, + r#"{"providers":{"meta":{"mechanism":"oauth","access_token":"dca:file-token"}}}"#, + ) + .unwrap(); + assert_eq!( + resolve_device_token_from(&HashMap::new(), dir.path()).unwrap(), + "dca:file-token" + ); + let env = environment(&[(MUSE_DEVICE_TOKEN_ENV, " dca:env-token ")]); + assert_eq!( + resolve_device_token_from(&env, dir.path()).unwrap(), + "dca:env-token" + ); + } + + #[test] + fn oauth_without_inline_token_does_not_fall_through_to_inference_key() { + let dir = tempdir().unwrap(); + let path = dir.path().join("auth.json"); + fs::write( + &path, + r#"{"providers":{"meta":{"mechanism":"oauth","api_key":"LLM-nope"}}}"#, + ) + .unwrap(); + let path = path.to_string_lossy().into_owned(); + let env = environment(&[(MUSE_AUTH_PATH_ENV, &path)]); + assert!(matches!( + resolve_device_token_from(&env, dir.path()), + Err(ProviderError::AuthRequired) + )); + } + + #[test] + fn success_payload_maps_windows_identity_details_and_ignores_secrets() { + let body = serde_json::to_vec(&success_payload()).unwrap(); + let result = parse_response(&body).unwrap(); + assert_eq!(result.source_label, "oauth"); + assert_eq!(result.usage.primary.used_percent, 96.0); + assert_eq!(result.usage.primary.window_minutes, Some(300)); + assert_eq!(result.usage.secondary.as_ref().unwrap().used_percent, 40.0); + assert_eq!( + result.usage.secondary.as_ref().unwrap().window_minutes, + Some(10080) + ); + assert_eq!( + result.usage.account_email.as_deref(), + Some("muse@example.com") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("Pro")); + let details: Vec<_> = result.display_details().collect(); + assert_eq!(details.len(), 3); + assert!( + details + .iter() + .any(|detail| detail.title() == "Plan" && detail.value() == "Pro") + ); + assert!( + details + .iter() + .all(|detail| !detail.value().contains("LLM_")) + ); + assert!( + details + .iter() + .all(|detail| !detail.value().contains("card_")) + ); + } + + #[test] + fn percentage_values_clamp_and_invalid_reset_omits_only_reset() { + let mut payload = success_payload(); + payload["subs_usage"]["window"]["used_percent"] = serde_json::json!(-5.0); + payload["subs_usage"]["weekly"]["used_percent"] = serde_json::json!(150.0); + payload["subs_usage"]["window"]["resets_at"] = serde_json::json!(64092211201_i64); + let result = parse_response(&serde_json::to_vec(&payload).unwrap()).unwrap(); + assert_eq!(result.usage.primary.used_percent, 0.0); + assert_eq!(result.usage.secondary.as_ref().unwrap().used_percent, 100.0); + assert!(result.usage.primary.resets_at.is_none()); + assert!(result.usage.secondary.as_ref().unwrap().resets_at.is_some()); + } + + #[test] + fn malformed_required_fields_fail_closed_without_echoing_payload() { + for payload in [ + serde_json::json!({"is_subs_active": true}), + serde_json::json!({"is_subs_active": true, "subs_usage": {"window": {}, "weekly": {}}}), + serde_json::json!({"require_payment": "yes", "is_subs_active": true}), + serde_json::json!({"require_payment": false, "is_subs_active": true, "subs_usage": {"window": {"window_duration_mins": "300"}, "weekly": {"used_percent": 1}}}), + ] { + let error = parse_response(&serde_json::to_vec(&payload).unwrap()).unwrap_err(); + assert!(matches!(error, ProviderError::Parse(_))); + assert!(!error.to_string().contains("LLM-nope")); + } + } + + #[test] + fn subscription_errors_are_safe_and_statuses_are_classified() { + let mut payload = success_payload(); + payload["require_payment"] = serde_json::json!(true); + let error = parse_response(&serde_json::to_vec(&payload).unwrap()).unwrap_err(); + assert!(error.to_string().contains("dev.meta.ai")); + assert!(!error.to_string().contains("LLM_")); + + assert!(matches!( + status_error(StatusCode::UNAUTHORIZED), + ProviderError::AuthRequired + )); + assert!(matches!( + status_error(StatusCode::FORBIDDEN), + ProviderError::AuthRequired + )); + assert!( + status_error(StatusCode::TOO_MANY_REQUESTS) + .to_string() + .contains("rate limited") + ); + assert!( + status_error(StatusCode::INTERNAL_SERVER_ERROR) + .to_string() + .contains("unavailable") + ); + } +} From 02903e6bfa0c6f394f2ee6c8f72b666e52cf6625 Mon Sep 17 00:00:00 2001 From: NessZerra <90105158+Finesssee@users.noreply.github.com> Date: Sun, 20 Sep 2026 01:46:14 +0700 Subject: [PATCH 2/2] Bound Muse response buffering --- rust/src/providers/muse/mod.rs | 50 ++++++++++++++++++++++++++++------ 1 file changed, 42 insertions(+), 8 deletions(-) diff --git a/rust/src/providers/muse/mod.rs b/rust/src/providers/muse/mod.rs index 83bdabbb53..fbe54d14a6 100644 --- a/rust/src/providers/muse/mod.rs +++ b/rust/src/providers/muse/mod.rs @@ -7,6 +7,7 @@ use async_trait::async_trait; use chrono::{DateTime, Utc}; +use futures::StreamExt; use reqwest::{Client, StatusCode}; use serde::Deserialize; use serde_json::Value; @@ -101,16 +102,31 @@ impl MuseProvider { )); } - let body = response.bytes().await.map_err(ProviderError::Network)?; - if body.len() > MAX_RESPONSE_BYTES { - return Err(ProviderError::Parse( - "Muse Code returned an oversized response.".to_string(), - )); - } + let body = read_bounded_body(response).await?; parse_response(&body) } } +async fn read_bounded_body(response: reqwest::Response) -> Result, ProviderError> { + 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(ProviderError::Parse( + "Muse Code returned an oversized response.".to_string(), + )); + } + body.extend_from_slice(chunk); + Ok(()) +} + impl Default for MuseProvider { fn default() -> Self { Self::new() @@ -358,14 +374,18 @@ fn optional_text(value: Option<&Value>, field: &str) -> Result, P } fn positive_safe_minutes(value: f64) -> Result { - if value <= 0.0 || value.fract() != 0.0 || value > u32::MAX as f64 { + if !value.is_finite() || value <= 0.0 { + return Err(parse_failure("window_duration_mins")); + } + let rounded = value.round(); + if rounded <= 0.0 || rounded > u32::MAX as f64 { return Err(parse_failure("window_duration_mins")); } #[allow( clippy::cast_possible_truncation, reason = "value is bounded to u32 range" )] - let value = value as u32; + let value = rounded as u32; Ok(value) } @@ -608,4 +628,18 @@ mod tests { .contains("unavailable") ); } + + #[test] + fn fractional_window_duration_rounds_safely() { + let mut payload = success_payload(); + payload["subs_usage"]["window"]["window_duration_mins"] = serde_json::json!(300.6); + let result = parse_response(&serde_json::to_vec(&payload).unwrap()).unwrap(); + assert_eq!(result.usage.primary.window_minutes, Some(301)); + } + + #[test] + fn streaming_response_cap_rejects_oversized_chunk_without_content_length() { + let mut body = vec![0_u8; MAX_RESPONSE_BYTES]; + assert!(append_bounded_body(&mut body, &[0]).is_err()); + } }