vtcode_config/models/
mimo_auth.rs

1use std::fmt;
2use std::str::FromStr;
3
4use serde::{Deserialize, Serialize};
5
6use crate::constants::{env_vars, urls};
7
8/// The env key used by Provider::MiMo for pay-as-you-go
9const MIMO_API_KEY: &str = "MIMO_API_KEY";
10
11/// Authentication method for Xiaomi MiMo provider
12#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))]
13#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
14#[serde(rename_all = "kebab-case")]
15pub enum MiMoAuthMethod {
16    /// Pay-as-you-go: uses api-key header, sk- prefix, api.xiaomimimo.com/v1
17    #[default]
18    #[serde(alias = "payg", alias = "pay_as_you_go", alias = "apikey", alias = "api_key")]
19    PayAsYouGo,
20    /// Token Plan: uses Authorization Bearer header, tp- prefix, token-plan-cn.xiaomimimo.com/v1
21    #[serde(alias = "token_plan", alias = "tokenplan", alias = "tp")]
22    TokenPlan,
23    /// Forward-compatible catch-all for unknown auth methods
24    #[serde(other)]
25    Unknown,
26}
27
28impl MiMoAuthMethod {
29    #[must_use]
30    pub fn label(&self) -> &'static str {
31        match self {
32            Self::PayAsYouGo => "API Usage Billing",
33            Self::TokenPlan => "Subscription Plan",
34            Self::Unknown => "Unknown",
35        }
36    }
37
38    #[must_use]
39    pub fn description(&self) -> &'static str {
40        match self {
41            Self::PayAsYouGo => "Standard API access. Uses sk- key with api-key header.",
42            Self::TokenPlan => {
43                "Subscription-based access. Uses tp- key with Bearer token. Includes more models. Defaults to Europe cluster."
44            }
45            Self::Unknown => "Unrecognized authentication method.",
46        }
47    }
48
49    #[must_use]
50    pub fn env_key(&self) -> &'static str {
51        match self {
52            Self::PayAsYouGo => MIMO_API_KEY,
53            Self::TokenPlan => env_vars::MIMO_TOKEN_PLAN_KEY,
54            Self::Unknown => MIMO_API_KEY,
55        }
56    }
57
58    #[must_use]
59    pub fn api_base(&self) -> &'static str {
60        match self {
61            Self::PayAsYouGo => urls::MIMO_API_BASE,
62            Self::TokenPlan => urls::MIMO_TOKEN_PLAN_API_BASE,
63            Self::Unknown => urls::MIMO_API_BASE,
64        }
65    }
66}
67
68impl fmt::Display for MiMoAuthMethod {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        match self {
71            Self::PayAsYouGo => write!(f, "api-usage-billing"),
72            Self::TokenPlan => write!(f, "subscription-plan"),
73            Self::Unknown => write!(f, "unknown"),
74        }
75    }
76}
77
78impl FromStr for MiMoAuthMethod {
79    type Err = String;
80
81    fn from_str(s: &str) -> Result<Self, Self::Err> {
82        match s.to_ascii_lowercase().as_str() {
83            "payg" | "pay-as-you-go" | "pay_as_you_go" | "apikey" | "api-key" | "api-usage-billing" => {
84                Ok(Self::PayAsYouGo)
85            }
86            "token-plan" | "token_plan" | "tokenplan" | "tp" | "subscription-plan" => Ok(Self::TokenPlan),
87            _ => Err(format!("Unknown MiMo auth method: {s}")),
88        }
89    }
90}
91
92/// Detect auth method from API key prefix or base URL.
93///
94/// Returns `PayAsYouGo` as the default if no hint is available.
95#[must_use]
96pub fn detect_mimo_auth_method(api_key: &str, base_url: Option<&str>) -> MiMoAuthMethod {
97    if api_key.starts_with("tp-") {
98        return MiMoAuthMethod::TokenPlan;
99    }
100    if api_key.starts_with("sk-") {
101        return MiMoAuthMethod::PayAsYouGo;
102    }
103    // If no key prefix hint, check base URL
104    if let Some(url) = base_url
105        && url.contains("token-plan")
106    {
107        return MiMoAuthMethod::TokenPlan;
108    }
109    // Check env var for base URL override
110    if let Ok(url) = std::env::var(env_vars::MIMO_TOKEN_PLAN_BASE_URL)
111        && !url.trim().is_empty()
112    {
113        return MiMoAuthMethod::TokenPlan;
114    }
115    MiMoAuthMethod::PayAsYouGo
116}
117
118#[cfg(test)]
119mod tests {
120    use super::*;
121
122    #[test]
123    fn detect_from_key_prefix() {
124        assert_eq!(detect_mimo_auth_method("sk-abc123", None), MiMoAuthMethod::PayAsYouGo);
125        assert_eq!(detect_mimo_auth_method("tp-abc123", None), MiMoAuthMethod::TokenPlan);
126    }
127
128    #[test]
129    fn detect_from_base_url() {
130        assert_eq!(
131            detect_mimo_auth_method("abc", Some("https://token-plan-cn.xiaomimimo.com/v1")),
132            MiMoAuthMethod::TokenPlan
133        );
134        assert_eq!(detect_mimo_auth_method("abc", Some("https://api.xiaomimimo.com/v1")), MiMoAuthMethod::PayAsYouGo);
135    }
136
137    #[test]
138    fn detect_defaults_to_payg() {
139        assert_eq!(detect_mimo_auth_method("", None), MiMoAuthMethod::PayAsYouGo);
140        assert_eq!(detect_mimo_auth_method("unknown", None), MiMoAuthMethod::PayAsYouGo);
141    }
142
143    #[test]
144    fn parse_from_str() {
145        assert_eq!("payg".parse::<MiMoAuthMethod>().unwrap(), MiMoAuthMethod::PayAsYouGo);
146        assert_eq!("token-plan".parse::<MiMoAuthMethod>().unwrap(), MiMoAuthMethod::TokenPlan);
147        assert_eq!("tp".parse::<MiMoAuthMethod>().unwrap(), MiMoAuthMethod::TokenPlan);
148        assert!("invalid".parse::<MiMoAuthMethod>().is_err());
149    }
150
151    #[test]
152    fn display_roundtrip() {
153        let payg = MiMoAuthMethod::PayAsYouGo;
154        assert_eq!(payg.to_string(), "api-usage-billing");
155        let tp = MiMoAuthMethod::TokenPlan;
156        assert_eq!(tp.to_string(), "subscription-plan");
157    }
158}
vtcode_config/models/mimo_auth.rs

vtcode_config/models/
mimo_auth.rs
