magi_code/config/

custom_provider_config.rs

use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use std::collections::{BTreeMap, BTreeSet};

use super::Settings;

#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
#[serde(rename_all = "kebab-case")]
pub enum CustomReasoningProtocol {
    #[default]
    GptLike,
    AnthropicLike,
}

pub(crate) const MAX_CUSTOM_PROVIDER_REQUEST_HEADERS: usize = 32;

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
#[serde(tag = "source", rename_all = "snake_case")]
pub enum CustomProviderHeaderValue {
    ConversationId,
}

#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema, PartialEq, Eq)]
pub struct CustomProviderConfig {
    pub label: String,
    pub base_url: String,
    #[serde(
        default,
        deserialize_with = "deserialize_optional_env_var",
        skip_serializing_if = "Option::is_none"
    )]
    pub api_key_env_var: Option<String>,
    #[serde(
        default,
        deserialize_with = "deserialize_optional_models_dev_provider",
        skip_serializing_if = "Option::is_none"
    )]
    pub models_dev_provider: Option<String>,
    #[serde(
        default,
        deserialize_with = "deserialize_optional_fast_mode",
        skip_serializing_if = "Option::is_none"
    )]
    pub fast_mode: Option<CustomProviderFastMode>,
    #[serde(default, skip_serializing_if = "is_false")]
    pub use_responses_endpoint: bool,
    #[serde(default, skip_serializing_if = "is_false")]
    pub supports_text_verbosity: bool,
    #[serde(default, skip_serializing_if = "is_gpt_like")]
    pub reasoning_protocol: CustomReasoningProtocol,
    #[serde(
        default,
        deserialize_with = "deserialize_extra_models",
        skip_serializing_if = "Vec::is_empty"
    )]
    pub extra_models: Vec<String>,
    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
    pub request_headers: BTreeMap<String, CustomProviderHeaderValue>,
}

pub(super) fn validate_custom_provider_settings(settings: &Settings) -> anyhow::Result<()> {
    for (id, custom) in &settings.custom_providers {
        validate_custom_provider_id(id).map_err(|error| {
            anyhow::anyhow!("custom provider '{id}' has invalid provider id: {error}")
        })?;
        validate_custom_provider_label(&custom.label).map_err(|error| {
            anyhow::anyhow!("custom provider '{id}' has invalid label: {error}")
        })?;
        normalize_custom_provider_base_url(&custom.base_url).map_err(|error| {
            anyhow::anyhow!("custom provider '{id}' has invalid base_url: {error}")
        })?;
        if let Some(env_var) = &custom.api_key_env_var {
            validate_env_var_name(env_var).map_err(|error| {
                anyhow::anyhow!("custom provider '{id}' has invalid api_key_env_var: {error}")
            })?;
        }
        if let Some(fast_mode) = &custom.fast_mode {
            validate_custom_provider_fast_mode(fast_mode).map_err(|error| {
                anyhow::anyhow!("custom provider '{id}' has invalid fast_mode: {error}")
            })?;
        }
        if let Some(models_dev_provider) = &custom.models_dev_provider {
            validate_models_dev_provider_namespace(models_dev_provider).map_err(|error| {
                anyhow::anyhow!("custom provider '{id}' has invalid models_dev_provider: {error}")
            })?;
        }
        normalized_extra_models(&custom.extra_models).map_err(|error| {
            anyhow::anyhow!("custom provider '{id}' has invalid extra_models: {error}")
        })?;
        validate_custom_provider_request_headers(id, &custom.request_headers)?;
    }
    Ok(())
}

fn validate_custom_provider_request_headers(
    provider_id: &str,
    headers: &BTreeMap<String, CustomProviderHeaderValue>,
) -> anyhow::Result<()> {
    if headers.len() > MAX_CUSTOM_PROVIDER_REQUEST_HEADERS {
        anyhow::bail!(
            "custom provider '{provider_id}' request_headers must contain at most {MAX_CUSTOM_PROVIDER_REQUEST_HEADERS} entries"
        );
    }
    let mut normalized = BTreeSet::new();
    for name in headers.keys() {
        reqwest::header::HeaderName::from_bytes(name.as_bytes()).map_err(|_| {
            anyhow::anyhow!(
                "custom provider '{provider_id}' has invalid request header name '{name}'"
            )
        })?;
        let lower = name.to_ascii_lowercase();
        if !normalized.insert(lower.clone()) {
            anyhow::bail!(
                "custom provider '{provider_id}' has duplicate request header name '{name}' (names are case-insensitive)"
            );
        }
        if matches!(
            lower.as_str(),
            "accept"
                | "authorization"
                | "content-length"
                | "content-type"
                | "host"
                | "proxy-authorization"
                | "transfer-encoding"
                | "user-agent"
        ) {
            anyhow::bail!(
                "custom provider '{provider_id}' request_headers must not override transport-owned header '{name}'"
            );
        }
    }
    Ok(())
}

fn validate_custom_provider_label(label: &str) -> anyhow::Result<()> {
    let label = label.trim();
    if label.is_empty() || label.len() > 100 {
        anyhow::bail!("custom provider label must be non-empty and at most 100 characters");
    }
    if looks_like_secret_label(label) {
        anyhow::bail!("custom provider label must not look like a secret value");
    }
    Ok(())
}

fn looks_like_secret_label(value: &str) -> bool {
    let value = value.trim();
    value.starts_with("sk-")
        || value.starts_with("Bearer ")
        || value.contains('=')
        || (value.len() >= 48
            && value
                .chars()
                .filter(|ch| ch.is_ascii_alphanumeric())
                .count()
                >= 40)
}

pub(crate) fn validate_custom_provider_id(id: &str) -> anyhow::Result<String> {
    let id = id.trim();
    if matches!(
        id,
        crate::providers::OPENAI_CODEX_PROVIDER
            | crate::providers::ANTHROPIC_PROVIDER
            | "claude-code"
            | "openai"
    ) {
        anyhow::bail!("custom provider id '{id}' is reserved");
    }
    if id.len() > 63
        || id.is_empty()
        || !id.as_bytes()[0].is_ascii_lowercase()
        || id.ends_with('-')
        || !id
            .chars()
            .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-')
    {
        anyhow::bail!(
            "custom provider id must match ^[a-z][a-z0-9-]{{0,62}}$ with no trailing hyphen"
        );
    }
    Ok(id.to_string())
}

pub(crate) fn looks_like_secret_value(value: &str) -> bool {
    let value = value.trim();
    value.starts_with("sk-")
        || value.starts_with("Bearer ")
        || value.contains('=')
        || value.chars().any(char::is_whitespace)
        || (value.len() >= 48
            && value
                .chars()
                .filter(|ch| ch.is_ascii_alphanumeric())
                .count()
                >= 40)
}

pub(crate) fn validate_env_var_name(name: &str) -> anyhow::Result<String> {
    let name = name.trim();
    if looks_like_secret_value(name) {
        anyhow::bail!(
            "API key environment variable name looks like a secret value; enter a variable name such as CUSTOM_PROVIDER_API_KEY"
        );
    }
    if name.is_empty()
        || !(name.as_bytes()[0].is_ascii_uppercase() || name.as_bytes()[0] == b'_')
        || !name
            .chars()
            .all(|ch| ch.is_ascii_uppercase() || ch.is_ascii_digit() || ch == '_')
    {
        anyhow::bail!("API key environment variable name must match ^[A-Z_][A-Z0-9_]*$");
    }
    Ok(name.to_string())
}

pub(crate) fn validate_optional_env_var_name(name: &str) -> anyhow::Result<Option<String>> {
    if name.trim().is_empty() {
        return Ok(None);
    }
    validate_env_var_name(name).map(Some)
}

pub(crate) fn normalized_extra_models(extra_models: &[String]) -> anyhow::Result<Vec<String>> {
    if extra_models.len() > 64 {
        anyhow::bail!("extra_models must contain at most 64 model ids");
    }
    let mut seen = BTreeSet::new();
    let mut normalized = Vec::new();
    for model in extra_models {
        let model = model.trim();
        if model.is_empty() {
            anyhow::bail!("extra_models entries must be non-empty");
        }
        if model.len() > 200 {
            anyhow::bail!("extra_models entries must be at most 200 bytes");
        }
        if model
            .chars()
            .any(|ch| ch.is_ascii_control() || ch.is_ascii_whitespace())
        {
            anyhow::bail!(
                "extra_models entries must not contain ASCII control characters or whitespace"
            );
        }
        if looks_like_secret_value(model) {
            anyhow::bail!("extra_models entries must not look like secret values");
        }
        if seen.insert(model.to_string()) {
            normalized.push(model.to_string());
        }
    }
    Ok(normalized)
}

pub(crate) fn validate_models_dev_provider_namespace(namespace: &str) -> anyhow::Result<String> {
    let namespace = namespace.trim();
    if looks_like_secret_value(namespace) {
        anyhow::bail!(
            "models.dev provider namespace looks like a secret value; enter a namespace such as openai"
        );
    }
    if namespace.len() > 63
        || namespace.is_empty()
        || !namespace.as_bytes()[0].is_ascii_lowercase()
        || namespace.ends_with('-')
        || !namespace
            .chars()
            .all(|ch| ch.is_ascii_lowercase() || ch.is_ascii_digit() || ch == '-')
    {
        anyhow::bail!(
            "models.dev provider namespace must match ^[a-z][a-z0-9-]{{0,62}}$ with no trailing hyphen"
        );
    }
    Ok(namespace.to_string())
}

pub(crate) fn derive_custom_provider_id(label: &str) -> anyhow::Result<String> {
    let mut id = String::new();
    let mut last_was_separator = false;
    for ch in label.trim().chars() {
        if ch.is_ascii_alphanumeric() {
            id.push(ch.to_ascii_lowercase());
            last_was_separator = false;
        } else if !last_was_separator && !id.is_empty() {
            id.push('-');
            last_was_separator = true;
        }
    }
    while id.ends_with('-') {
        id.pop();
    }
    validate_custom_provider_id(&id)
        .map_err(|_| anyhow::anyhow!("custom provider label must derive a provider id matching ^[a-z][a-z0-9-]{{0,62}}$ and must not be reserved"))
}

pub(crate) fn normalize_custom_provider_base_url(input: &str) -> anyhow::Result<String> {
    let value = input.trim().trim_end_matches('/');
    let parsed = reqwest::Url::parse(value)
        .map_err(|_| anyhow::anyhow!("custom provider base URL must be a valid URL"))?;
    if !parsed.username().is_empty() || parsed.password().is_some() {
        anyhow::bail!("custom provider base URL must not include URL credentials or userinfo");
    }
    if parsed.query().is_some() || parsed.fragment().is_some() {
        anyhow::bail!("custom provider base URL must not include query parameters or fragments");
    }
    let path = parsed.path().trim_end_matches('/');
    if path.ends_with("/responses")
        || path.ends_with("/models")
        || path.ends_with("/completions")
        || path.ends_with("/chat/completions")
    {
        anyhow::bail!("custom provider base URL must be an API root, not an endpoint URL");
    }
    match parsed.scheme() {
        "https" | "http" => Ok(value.to_string()),
        _ => anyhow::bail!("custom provider base URL must use http:// or https://"),
    }
}

pub(crate) fn make_custom_provider_config(
    label: &str,
    base_url: &str,
    api_key_env_var: &str,
) -> anyhow::Result<CustomProviderConfig> {
    let label = label.trim();
    validate_custom_provider_label(label)?;
    Ok(CustomProviderConfig {
        label: label.to_string(),
        base_url: normalize_custom_provider_base_url(base_url)?,
        api_key_env_var: validate_optional_env_var_name(api_key_env_var)?,
        models_dev_provider: None,
        fast_mode: None,
        use_responses_endpoint: false,
        supports_text_verbosity: false,
        reasoning_protocol: CustomReasoningProtocol::default(),
        extra_models: Vec::new(),
        request_headers: BTreeMap::new(),
    })
}

fn validate_custom_provider_fast_mode(fast_mode: &CustomProviderFastMode) -> anyhow::Result<()> {
    validate_fast_service_tier(&fast_mode.service_tier)?;
    validate_fast_models(&fast_mode.models)?;
    Ok(())
}

fn validate_fast_service_tier(service_tier: &str) -> anyhow::Result<String> {
    let service_tier = service_tier.trim();
    if service_tier.is_empty()
        || service_tier.len() > 64
        || !service_tier.bytes().all(|byte| {
            byte.is_ascii_alphanumeric() || byte == b'_' || byte == b'-' || byte == b'.'
        })
    {
        anyhow::bail!(
            "fast_mode.service_tier must be a non-empty ASCII identifier of at most 64 characters"
        );
    }
    if looks_like_secret_value(service_tier) {
        anyhow::bail!("fast_mode.service_tier must not look like a secret value");
    }
    Ok(service_tier.to_string())
}

fn validate_fast_models(models: &[String]) -> anyhow::Result<()> {
    if models.is_empty() || models.len() > 64 {
        anyhow::bail!("fast_mode.models must contain 1 to 64 model ids");
    }
    if models.iter().any(|model| model == "*") && models.len() != 1 {
        anyhow::bail!("fast_mode.models wildcard must be the sole member");
    }
    let mut seen = BTreeSet::new();
    for model in models {
        if model.is_empty()
            || model.chars().count() > 200
            || model
                .chars()
                .any(|ch| ch.is_whitespace() || ch.is_control())
        {
            anyhow::bail!(
                "fast_mode.models entries must be non-empty model ids of at most 200 characters without whitespace or control characters"
            );
        }
        if model != "*" && looks_like_secret_value(model) {
            anyhow::bail!("fast_mode.models entries must not look like secret values");
        }
        if !seen.insert(model.as_str()) {
            anyhow::bail!("fast_mode.models must not contain duplicate model ids");
        }
    }
    Ok(())
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct CustomProviderFastMode {
    pub service_tier: String,
    pub models: Vec<String>,
}

#[derive(JsonSchema)]
#[allow(dead_code)]
struct CustomProviderFastModeSchema {
    #[schemars(regex(pattern = r"^\s*[A-Za-z0-9_.-]{1,64}\s*$"))]
    service_tier: String,
    #[schemars(
        length(min = 1, max = 64),
        inner(regex(pattern = r"^\S{1,200}$")),
        transform = add_fast_models_schema_constraints
    )]
    models: Vec<String>,
}

fn add_fast_models_schema_constraints(schema: &mut schemars::Schema) {
    let object = schema.ensure_object();
    object.insert("uniqueItems".to_string(), serde_json::json!(true));
    object.insert(
        "oneOf".to_string(),
        serde_json::json!([
            {
                "contains": {"pattern": r"^\*$"},
                "maxItems": 1
            },
            {
                "not": {"contains": {"pattern": r"^\*$"}}
            }
        ]),
    );
}

impl JsonSchema for CustomProviderFastMode {
    fn schema_name() -> std::borrow::Cow<'static, str> {
        "CustomProviderFastMode".into()
    }

    fn json_schema(generator: &mut schemars::SchemaGenerator) -> schemars::Schema {
        CustomProviderFastModeSchema::json_schema(generator)
    }
}

#[cfg(test)]
impl CustomProviderFastMode {
    pub(crate) fn supports_model(&self, model: &str) -> bool {
        self.models
            .iter()
            .any(|candidate| candidate == "*" || candidate == model)
    }
}
fn is_gpt_like(value: &CustomReasoningProtocol) -> bool {
    *value == CustomReasoningProtocol::GptLike
}

fn is_false(value: &bool) -> bool {
    !*value
}

fn deserialize_optional_env_var<'de, D>(deserializer: D) -> Result<Option<String>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<String>::deserialize(deserializer)?;
    Ok(value.and_then(|value| {
        let trimmed = value.trim();
        if trimmed.is_empty() {
            None
        } else {
            Some(trimmed.to_string())
        }
    }))
}

fn deserialize_optional_fast_mode<'de, D>(
    deserializer: D,
) -> Result<Option<CustomProviderFastMode>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<CustomProviderFastMode>::deserialize(deserializer)?;
    Ok(value.map(|mut fast| {
        fast.service_tier = fast.service_tier.trim().to_string();
        fast
    }))
}

fn deserialize_optional_models_dev_provider<'de, D>(
    deserializer: D,
) -> Result<Option<String>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let value = Option::<String>::deserialize(deserializer)?;
    Ok(value.and_then(|value| {
        let trimmed = value.trim();
        if trimmed.is_empty() {
            None
        } else {
            Some(trimmed.to_string())
        }
    }))
}

fn deserialize_extra_models<'de, D>(deserializer: D) -> Result<Vec<String>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let values = Vec::<String>::deserialize(deserializer)?;
    Ok(values
        .into_iter()
        .map(|value| value.trim().to_string())
        .collect())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn settings_with_fast_models(models: Vec<String>) -> Settings {
        Settings {
            custom_providers: [(
                "provider".to_string(),
                CustomProviderConfig {
                    label: "Provider".to_string(),
                    base_url: "https://example.test".to_string(),
                    api_key_env_var: None,
                    models_dev_provider: None,
                    fast_mode: Some(CustomProviderFastMode {
                        service_tier: "priority".to_string(),
                        models,
                    }),
                    use_responses_endpoint: false,
                    supports_text_verbosity: false,
                    reasoning_protocol: CustomReasoningProtocol::default(),
                    extra_models: Vec::new(),
                    request_headers: BTreeMap::new(),
                },
            )]
            .into(),
            ..Settings::default()
        }
    }

    #[test]
    fn fast_mode_trims_service_tier_but_preserves_exact_model_ids() {
        let config: CustomProviderConfig = serde_json::from_value(serde_json::json!({
            "label": "Provider", "base_url": "https://example.test",
            "fast_mode": {"service_tier": " priority ", "models": [" model-a "]}
        }))
        .unwrap();
        let fast_mode = config.fast_mode.as_ref().unwrap();
        assert_eq!(fast_mode.service_tier, "priority");
        assert_eq!(fast_mode.models, vec![" model-a ".to_string()]);
        assert!(!fast_mode.supports_model("model-a"));
        assert!(fast_mode.supports_model(" model-a "));
        assert!(
            validate_custom_provider_settings(&settings_with_fast_models(fast_mode.models.clone()))
                .is_err()
        );
    }

    #[test]
    fn request_headers_parse_conversation_source_and_reject_owned_or_duplicate_names() {
        let config: CustomProviderConfig = serde_json::from_value(serde_json::json!({
            "label": "Provider",
            "base_url": "https://example.test",
            "request_headers": {
                "x-opencode-session": {"source": "conversation_id"}
            }
        }))
        .unwrap();
        assert_eq!(
            config.request_headers.get("x-opencode-session"),
            Some(&CustomProviderHeaderValue::ConversationId)
        );

        for headers in [
            BTreeMap::from([(
                "Authorization".to_string(),
                CustomProviderHeaderValue::ConversationId,
            )]),
            BTreeMap::from([
                (
                    "X-Session".to_string(),
                    CustomProviderHeaderValue::ConversationId,
                ),
                (
                    "x-session".to_string(),
                    CustomProviderHeaderValue::ConversationId,
                ),
            ]),
        ] {
            assert!(validate_custom_provider_request_headers("provider", &headers).is_err());
        }
    }

    #[test]
    fn fast_mode_runtime_validation_uses_unicode_character_limits_and_rejects_whitespace() {
        for models in [
            vec![" model".to_string()],
            vec!["model ".to_string()],
            vec!["model id".to_string()],
            vec!["model\u{2003}id".to_string()],
            vec!["model\u{0000}id".to_string()],
        ] {
            assert!(validate_custom_provider_settings(&settings_with_fast_models(models)).is_err());
        }
        assert!(
            validate_custom_provider_settings(&settings_with_fast_models(vec!["😀".repeat(200),]))
                .is_ok()
        );
        assert!(
            validate_custom_provider_settings(&settings_with_fast_models(vec!["😀".repeat(201),]))
                .is_err()
        );
    }

    #[test]
    fn fast_mode_rejects_duplicate_models_and_non_sole_wildcard() {
        for models in [
            vec!["model".to_string(), "model".to_string()],
            vec!["*".to_string(), "model".to_string()],
        ] {
            assert!(validate_custom_provider_settings(&settings_with_fast_models(models)).is_err());
        }
    }
}