yoagent 0.26.5

//! Provider registry — maps ApiProtocol to StreamProvider implementations.

use super::model::{ApiProtocol, ModelConfig};
use super::traits::*;
use crate::types::*;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::mpsc;

/// Registry of all available stream providers, keyed by API protocol.
pub struct ProviderRegistry {
    providers: HashMap<ApiProtocol, Arc<dyn StreamProvider>>,
}

impl ProviderRegistry {
    /// Create an empty registry.
    pub fn new() -> Self {
        Self {
            providers: HashMap::new(),
        }
    }

    /// Register a provider for a given protocol.
    pub fn register(&mut self, protocol: ApiProtocol, provider: impl StreamProvider + 'static) {
        self.providers.insert(protocol, Arc::new(provider));
    }

    /// Resolve the provider for a protocol as a shared handle.
    ///
    /// Unlike [`get`](Self::get), this hands out an owned `Arc` that can be
    /// stored (e.g. by [`Agent::from_config`](crate::Agent::from_config)),
    /// so callers don't have to borrow the registry for the provider's
    /// lifetime.
    pub fn resolve(&self, protocol: &ApiProtocol) -> Option<Arc<dyn StreamProvider>> {
        self.providers.get(protocol).cloned()
    }

    /// Borrow the provider for a protocol for a transient call. Prefer
    /// [`resolve`](Self::resolve) when you need to keep the handle.
    pub fn get(&self, protocol: &ApiProtocol) -> Option<&dyn StreamProvider> {
        self.providers.get(protocol).map(|p| p.as_ref())
    }

    /// Check if a protocol is registered.
    pub fn has(&self, protocol: &ApiProtocol) -> bool {
        self.providers.contains_key(protocol)
    }

    /// List all registered protocols.
    pub fn protocols(&self) -> Vec<ApiProtocol> {
        self.providers.keys().copied().collect()
    }

    /// Stream using the appropriate provider for the model's API protocol.
    pub async fn stream(
        &self,
        model: &ModelConfig,
        config: StreamConfig,
        tx: mpsc::UnboundedSender<StreamEvent>,
        cancel: tokio_util::sync::CancellationToken,
    ) -> Result<Message, ProviderError> {
        let provider = self.providers.get(&model.api).ok_or_else(|| {
            ProviderError::Other(format!(
                "No provider registered for protocol: {}",
                model.api
            ))
        })?;

        provider.stream(config, tx, cancel).await
    }
}

impl Default for ProviderRegistry {
    /// Create a registry with all built-in providers registered.
    fn default() -> Self {
        use crate::provider::{
            AnthropicProvider, AzureOpenAiProvider, BedrockProvider, GoogleProvider,
            GoogleVertexProvider, OpenAiCompatProvider, OpenAiResponsesProvider,
        };

        let mut registry = Self::new();
        registry.register(ApiProtocol::AnthropicMessages, AnthropicProvider);
        registry.register(ApiProtocol::OpenAiCompletions, OpenAiCompatProvider);
        registry.register(ApiProtocol::OpenAiResponses, OpenAiResponsesProvider);
        registry.register(ApiProtocol::GoogleGenerativeAi, GoogleProvider);
        registry.register(ApiProtocol::GoogleVertex, GoogleVertexProvider);
        registry.register(ApiProtocol::BedrockConverseStream, BedrockProvider);
        registry.register(ApiProtocol::AzureOpenAiResponses, AzureOpenAiProvider);

        registry
    }
}

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

    #[test]
    fn test_default_registry_has_all_providers() {
        let registry = ProviderRegistry::default();

        assert!(registry.has(&ApiProtocol::AnthropicMessages));
        assert!(registry.has(&ApiProtocol::OpenAiCompletions));
        assert!(registry.has(&ApiProtocol::OpenAiResponses));
        assert!(registry.has(&ApiProtocol::GoogleGenerativeAi));
        assert!(registry.has(&ApiProtocol::GoogleVertex));
        assert!(registry.has(&ApiProtocol::BedrockConverseStream));
        assert!(registry.has(&ApiProtocol::AzureOpenAiResponses));
    }

    #[test]
    fn test_registry_protocols() {
        let registry = ProviderRegistry::default();
        let protocols = registry.protocols();
        assert_eq!(protocols.len(), 7);
    }

    #[test]
    fn test_custom_registry() {
        let mut registry = ProviderRegistry::new();
        assert!(!registry.has(&ApiProtocol::AnthropicMessages));

        registry.register(
            ApiProtocol::AnthropicMessages,
            crate::provider::AnthropicProvider,
        );
        assert!(registry.has(&ApiProtocol::AnthropicMessages));
    }

    #[test]
    fn api_key_variables_lists_every_credential_once() {
        let names = super::api_key_variables();
        for name in [
            "ANTHROPIC_API_KEY",
            "GOOGLE_API_KEY",
            "MODEL_API_KEY",
            "OPENCODE_API_KEY",
            "AZURE_OPENAI_API_KEY",
            "AWS_BEARER_TOKEN_BEDROCK",
            "AWS_SECRET_ACCESS_KEY",
            "AWS_SESSION_TOKEN",
            "YOAGENT_API_KEY",
            "API_KEY",
            "TYPESAFE_API_KEY",
            "CLOUDFLARE_AUTH_TOKEN",
        ] {
            assert!(names.contains(&name), "{name} missing from {names:?}");
        }
        let mut unique = names.clone();
        unique.sort_unstable();
        unique.dedup();
        assert_eq!(unique.len(), names.len(), "a name listed twice: {names:?}");
        // Every provider's own variables come first, in the table's order.
        assert_eq!(&names[..2], ["ANTHROPIC_API_KEY", "OPENAI_API_KEY"]);
    }

    /// The decision presets read these names from their own constants.
    #[cfg(feature = "decision")]
    #[test]
    fn decision_key_variables_match_the_presets() {
        use crate::decision::{
            CLOUDFLARE_API_TOKEN_ENV, CLOUDFLARE_AUTH_TOKEN_ENV, TYPESAFE_API_KEY_ENV,
        };
        assert_eq!(
            super::DECISION_KEY_VARIABLES,
            [
                TYPESAFE_API_KEY_ENV,
                CLOUDFLARE_API_TOKEN_ENV,
                CLOUDFLARE_AUTH_TOKEN_ENV
            ]
        );
    }
}

/// The variables [`resolve_api_key`] reads for each provider name with keys of
/// its own, first match wins (Bedrock's, `vertex` and the keyless locals are
/// handled apart; any other name reads [`FALLBACK_KEY_VARIABLES`]).
const PROVIDER_KEY_VARIABLES: &[(&str, &[&str])] = &[
    ("anthropic", &["ANTHROPIC_API_KEY"]),
    ("openai", &["OPENAI_API_KEY"]),
    ("google", &["GEMINI_API_KEY", "GOOGLE_API_KEY"]),
    ("xai", &["XAI_API_KEY"]),
    ("groq", &["GROQ_API_KEY"]),
    ("deepseek", &["DEEPSEEK_API_KEY"]),
    ("mistral", &["MISTRAL_API_KEY"]),
    ("zai", &["ZAI_API_KEY"]),
    ("minimax", &["MINIMAX_API_KEY"]),
    // Meta's docs name the generic MODEL_API_KEY; prefer the unambiguous
    // META_API_KEY, accept the official one.
    ("meta", &["META_API_KEY", "MODEL_API_KEY"]),
    ("openrouter", &["OPENROUTER_API_KEY"]),
    ("cerebras", &["CEREBRAS_API_KEY"]),
    ("qwen", &["DASHSCOPE_API_KEY"]),
    ("opencode-zen", &["OPENCODE_API_KEY"]),
    ("opencode-go", &["OPENCODE_API_KEY"]),
    ("azure", &["AZURE_OPENAI_API_KEY"]),
];

/// What [`resolve_api_key`] reads for a provider name it does not know.
const FALLBACK_KEY_VARIABLES: &[&str] = &["YOAGENT_API_KEY", "API_KEY"];

/// Bedrock's credentials: a Bedrock API key, else an AWS access key pair (and
/// session token).
const BEDROCK_KEY_VARIABLES: &[&str] = &[
    crate::provider::bedrock::BEARER_TOKEN_ENV,
    "AWS_ACCESS_KEY_ID",
    "AWS_SECRET_ACCESS_KEY",
    "AWS_SESSION_TOKEN",
];

/// Credentials the decision-model presets read (feature `decision`), beyond
/// the provider keys above: TypeSafe's (`jev()`) and Cloudflare's (`clef*()`).
const DECISION_KEY_VARIABLES: &[&str] = &[
    "TYPESAFE_API_KEY",
    "CLOUDFLARE_API_TOKEN",
    "CLOUDFLARE_AUTH_TOKEN",
];

/// Every environment variable yoagent reads a credential from: each
/// provider's key variables as [`resolve_api_key`] reads them (Bedrock's
/// AWS ones included), its `YOAGENT_API_KEY` / `API_KEY` fallback, and the
/// decision-model presets' keys — each once, in that order.
///
/// For hosts that keep credentials away from code they run (plugins, MCP
/// servers): set each of these to the empty string in the child's
/// environment when the launcher can only add variables. A key a host
/// passes under another name (a manifest's own variable) is not in it.
pub fn api_key_variables() -> Vec<&'static str> {
    let mut names: Vec<&'static str> = Vec::new();
    let all = PROVIDER_KEY_VARIABLES
        .iter()
        .flat_map(|(_, vars)| vars.iter())
        .chain(BEDROCK_KEY_VARIABLES)
        .chain(FALLBACK_KEY_VARIABLES)
        .chain(DECISION_KEY_VARIABLES);
    for name in all {
        if !names.contains(name) {
            names.push(name);
        }
    }
    names
}

/// Resolve an API key from the conventional environment variable(s) for a
/// provider name (the `ModelConfig.provider` string).
///
/// Used as a fallback when no explicit key is set on the agent — an explicit
/// `with_api_key` always wins. Conventions:
///
/// | provider | env var(s), first match wins |
/// |---|---|
/// | `anthropic` | `ANTHROPIC_API_KEY` |
/// | `openai` | `OPENAI_API_KEY` |
/// | `google` | `GEMINI_API_KEY`, `GOOGLE_API_KEY` |
/// | `xai` / `groq` / `deepseek` / `mistral` / `zai` / `minimax` / `openrouter` / `cerebras` | `<PROVIDER>_API_KEY` |
/// | `qwen` | `DASHSCOPE_API_KEY` |
/// | `meta` | `META_API_KEY`, then `MODEL_API_KEY` |
/// | `opencode-zen` / `opencode-go` | `OPENCODE_API_KEY` |
/// | `azure` | `AZURE_OPENAI_API_KEY` |
/// | `bedrock` | `AWS_BEARER_TOKEN_BEDROCK` (a Bedrock API key, sent as a bearer token), then `AWS_ACCESS_KEY_ID` + `AWS_SECRET_ACCESS_KEY` (+ `AWS_SESSION_TOKEN`), composed as `access:secret[:token]` and used to SigV4-sign. `Agent`, `SubAgentTool` and `LlmCompaction` do not call this for a `BedrockConverseStream` config, whatever its provider string: they leave the key empty and the Bedrock provider reads these variables per request |
/// | `vertex` | none — pass a short-lived OAuth token via `with_api_key` |
/// | `local` / `ollama` | no key needed (empty) |
/// | anything else | `YOAGENT_API_KEY`, then `API_KEY` |
pub fn resolve_api_key(provider: &str) -> Option<String> {
    use std::env::var;
    let first = |names: &[&str]| {
        names.iter().find_map(|n| {
            var(n).ok().inspect(|_| {
                tracing::debug!("resolved API key for provider '{}' from ${}", provider, n)
            })
        })
    };
    match provider {
        "bedrock" => {
            // Empty values count as unset (a blanked-out export).
            let set = |n: &str| var(n).ok().filter(|v| !v.trim().is_empty());
            // A Bedrock API key wins, as in the AWS SDKs.
            if let Some(token) = set(crate::provider::bedrock::BEARER_TOKEN_ENV) {
                tracing::debug!(
                    "resolved API key for provider 'bedrock' from ${}",
                    crate::provider::bedrock::BEARER_TOKEN_ENV
                );
                return Some(token);
            }
            let [_, access, secret, token] = BEDROCK_KEY_VARIABLES else {
                unreachable!("BEDROCK_KEY_VARIABLES lists four names")
            };
            let access = set(access)?;
            let secret = set(secret)?;
            Some(match set(token) {
                Some(token) => format!("{}:{}:{}", access, secret, token),
                None => format!("{}:{}", access, secret),
            })
        }
        "vertex" => None,
        "local" | "ollama" => Some(String::new()),
        other => first(
            PROVIDER_KEY_VARIABLES
                .iter()
                .find(|(name, _)| *name == other)
                .map_or(FALLBACK_KEY_VARIABLES, |(_, vars)| vars),
        ),
    }
}

/// Like [`resolve_api_key`], but warns (via `tracing`) when resolution fails
/// for a provider that needs a key, then falls back to an empty string so
/// the provider returns a clear authentication error instead of the failure
/// being invisible until the first request.
///
/// Resolution follows the config's **protocol** before its provider string:
/// a `BedrockConverseStream` config always gets an empty key, whatever its
/// provider is called (`ModelConfig::custom(.., "aws-bedrock", ..)` must not
/// pick up `API_KEY`), and the Bedrock provider reads
/// `AWS_BEARER_TOKEN_BEDROCK` / `AWS_*` itself when the request is built — so
/// it can say exactly which variable is missing or malformed. The warning
/// is still logged here when the environment has no Bedrock credentials.
/// Without a config the provider is taken to be `anthropic`.
pub(crate) fn resolve_api_key_or_warn(config: Option<&ModelConfig>) -> String {
    if let Some(config) = config {
        if config.api == ApiProtocol::BedrockConverseStream {
            if resolve_api_key("bedrock").is_none() {
                tracing::warn!(
                    "no API key found for provider '{}' (Amazon Bedrock): {}; requests \
                     will fail with an authentication error",
                    config.provider,
                    api_key_env_hint("bedrock")
                );
            }
            return String::new();
        }
    }
    let provider = config.map(|c| c.provider.as_str()).unwrap_or("anthropic");
    match resolve_api_key(provider) {
        Some(key) => key,
        None => {
            tracing::warn!(
                "no API key found for provider '{}': {}; requests will fail \
                 with an authentication error",
                provider,
                api_key_env_hint(provider)
            );
            String::new()
        }
    }
}

/// Remedy fragment for the missing-key warning, matching the
/// [`resolve_api_key`] table.
fn api_key_env_hint(provider: &str) -> &'static str {
    match provider {
        "anthropic" => "set ANTHROPIC_API_KEY or call .with_api_key(...)",
        "openai" => "set OPENAI_API_KEY or call .with_api_key(...)",
        "google" => "set GEMINI_API_KEY (or GOOGLE_API_KEY) or call .with_api_key(...)",
        "xai" => "set XAI_API_KEY or call .with_api_key(...)",
        "groq" => "set GROQ_API_KEY or call .with_api_key(...)",
        "deepseek" => "set DEEPSEEK_API_KEY or call .with_api_key(...)",
        "mistral" => "set MISTRAL_API_KEY or call .with_api_key(...)",
        "zai" => "set ZAI_API_KEY or call .with_api_key(...)",
        "minimax" => "set MINIMAX_API_KEY or call .with_api_key(...)",
        "meta" => "set META_API_KEY (or MODEL_API_KEY) or call .with_api_key(...)",
        "openrouter" => "set OPENROUTER_API_KEY or call .with_api_key(...)",
        "cerebras" => "set CEREBRAS_API_KEY or call .with_api_key(...)",
        "qwen" => "set DASHSCOPE_API_KEY or call .with_api_key(...)",
        "opencode-zen" | "opencode-go" => "set OPENCODE_API_KEY or call .with_api_key(...)",
        "azure" => "set AZURE_OPENAI_API_KEY or call .with_api_key(...)",
        "bedrock" => {
            "set AWS_BEARER_TOKEN_BEDROCK, or AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY \
             (+ AWS_SESSION_TOKEN), or call .with_api_key(...) with a Bedrock API key \
             or \"access:secret[:token]\""
        }
        "vertex" => "pass a short-lived OAuth token via .with_api_key(...)",
        _ => "set YOAGENT_API_KEY (or API_KEY) or call .with_api_key(...)",
    }
}

#[cfg(test)]
mod resolve_key_tests {
    use super::resolve_api_key;

    #[test]
    fn test_deterministic_branches() {
        // local/ollama need no key
        assert_eq!(resolve_api_key("local").as_deref(), Some(""));
        assert_eq!(resolve_api_key("ollama").as_deref(), Some(""));
        // vertex expects an explicit OAuth token — never env-resolved
        assert_eq!(resolve_api_key("vertex"), None);
    }

    #[test]
    fn test_meta_env_resolution() {
        // Own vars as of this writing — grep before reusing MODEL_API_KEY in
        // another test. Clear first: a developer shell may export these.
        std::env::remove_var("META_API_KEY");
        std::env::set_var("MODEL_API_KEY", "official-var");
        assert_eq!(resolve_api_key("meta").as_deref(), Some("official-var"));
        // The unambiguous var wins when both are set.
        std::env::set_var("META_API_KEY", "preferred-var");
        assert_eq!(resolve_api_key("meta").as_deref(), Some("preferred-var"));
        std::env::remove_var("META_API_KEY");
        std::env::remove_var("MODEL_API_KEY");
    }

    #[test]
    fn test_env_resolution() {
        // Use a provider name unique to this test to avoid env races with
        // parallel tests.
        std::env::set_var("YOAGENT_API_KEY", "from-generic-fallback");
        assert_eq!(
            resolve_api_key("some-unknown-gateway").as_deref(),
            Some("from-generic-fallback")
        );
        std::env::remove_var("YOAGENT_API_KEY");
    }
}

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

    /// Every `ApiProtocol` variant. The wildcard-free `match` is a compile
    /// forcing-function: adding a variant fails to build here until it's added
    /// to the list (and, in turn, registered in `default()` below), so the
    /// "default registry is complete" invariant behind `from_config`'s unwrap
    /// cannot silently rot.
    fn all_protocols() -> Vec<ApiProtocol> {
        fn _assert_exhaustive(p: ApiProtocol) {
            match p {
                ApiProtocol::AnthropicMessages
                | ApiProtocol::OpenAiCompletions
                | ApiProtocol::OpenAiResponses
                | ApiProtocol::AzureOpenAiResponses
                | ApiProtocol::GoogleGenerativeAi
                | ApiProtocol::GoogleVertex
                | ApiProtocol::BedrockConverseStream => {}
            }
        }
        vec![
            ApiProtocol::AnthropicMessages,
            ApiProtocol::OpenAiCompletions,
            ApiProtocol::OpenAiResponses,
            ApiProtocol::AzureOpenAiResponses,
            ApiProtocol::GoogleGenerativeAi,
            ApiProtocol::GoogleVertex,
            ApiProtocol::BedrockConverseStream,
        ]
    }

    /// The default registry must resolve a provider for every protocol, so
    /// `Agent::from_config` (which unwraps against it) can never panic.
    #[test]
    fn default_registry_covers_all_protocols() {
        let registry = ProviderRegistry::default();
        for api in all_protocols() {
            assert!(
                registry.resolve(&api).is_some(),
                "default registry missing a provider for {api}"
            );
        }
    }

    /// Each protocol must resolve to a provider that actually *speaks* that
    /// protocol — guards against a mis-wired `register()` line (e.g. mapping
    /// `GoogleVertex` to the Anthropic provider) that a mere is-some check
    /// would miss.
    #[test]
    fn default_registry_maps_each_protocol_to_its_own_provider() {
        let registry = ProviderRegistry::default();
        for api in all_protocols() {
            let provider = registry.resolve(&api).expect("registered");
            assert_eq!(
                provider.protocol(),
                Some(api),
                "protocol {api} resolved to a provider that reports {:?}",
                provider.protocol()
            );
        }
    }

    #[test]
    fn empty_registry_resolves_nothing() {
        assert!(ProviderRegistry::new()
            .resolve(&ApiProtocol::AnthropicMessages)
            .is_none());
    }
}