tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Add Hugging Face OAuth support, add auth tab to settings (#9552)

Browse Source 
Signed-off-by: jh-block <jhugo@block.xyz>
This commit is contained in:
jh-block
2026-06-03 15:30:33 +02:00
committed by GitHub
parent 9626b4c3c1
commit 30034b9b32
28 changed files with 3894 additions and 75 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/cli.rs
+6 -1
View File
@@ -1944,11 +1944,16 @@ async fn handle_local_models_command(command: LocalModelsCommand) -> Result<()>
// Download
let manager = goose::download_manager::get_download_manager();
let hf_token = goose::providers::huggingface_auth::resolve_token_async()
.await
.ok()
.flatten();
manager
.download_model_sharded(
.download_model_sharded_with_bearer_token(
format!("{}-model", model_id),
download_files,
file.size_bytes + mmproj_size_bytes,
hf_token,
None,
)
.await?;
crates/goose-server/src/openapi.rs
+6
View File
@@ -396,6 +396,8 @@ derive_utoipa!(IconTheme as IconThemeSchema);
super::routes::config_management::remove_extension,
super::routes::config_management::get_extensions,
super::routes::config_management::read_all_config,
super::routes::config_management::list_provider_secrets,
super::routes::config_management::delete_provider_secret,
super::routes::config_management::providers,
super::routes::config_management::get_provider_models,
super::routes::config_management::get_provider_model_info,
@@ -488,6 +490,10 @@ derive_utoipa!(IconTheme as IconThemeSchema);
super::routes::config_management::ConfigResponse,
super::routes::config_management::ProvidersResponse,
super::routes::config_management::ProviderDetails,
super::routes::config_management::ProviderSecretsResponse,
super::routes::config_management::ProviderSecret,
super::routes::config_management::ProviderSecretStorage,
super::routes::config_management::ProviderSecretStatus,
super::routes::config_management::SlashCommandsResponse,
super::routes::config_management::SlashCommand,
super::routes::config_management::CommandType,
crates/goose-server/src/routes/config_management.rs
+792 -25
View File
@@ -7,6 +7,7 @@ use axum::{
routing::{delete, get, post},
Json, Router,
};
use chrono::{DateTime, TimeZone, Utc};
use goose::config::declarative_providers::LoadedProvider;
use goose::config::paths::Paths;
use goose::config::ExtensionEntry;
@@ -20,6 +21,7 @@ use goose::providers::catalog::{
ProviderTemplate,
};
use goose::providers::create_with_default_model;
use goose::providers::huggingface_auth;
use goose::providers::providers as get_providers;
use goose::{
agents::execute_commands, agents::ExtensionConfig, config::permission::PermissionLevel,
@@ -28,7 +30,10 @@ use goose::{
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_yaml;
use std::{collections::HashMap, sync::Arc};
use std::{
collections::{HashMap, HashSet},
sync::Arc,
};
use utoipa::ToSchema;
#[derive(Serialize, ToSchema)]
@@ -141,6 +146,43 @@ pub enum ConfigValueResponse {
MaskedValue(MaskedSecret),
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ProviderSecretStorage {
SecretStore,
ProviderCache,
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ProviderSecretStatus {
Valid,
Expired,
Unknown,
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub struct ProviderSecret {
pub id: String,
pub provider: String,
pub provider_display_name: String,
pub name: String,
pub storage: ProviderSecretStorage,
pub expires_at: Option<DateTime<Utc>>,
pub status: ProviderSecretStatus,
pub configured: bool,
pub has_secret: bool,
pub can_delete: bool,
pub can_configure: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub configure_provider: Option<String>,
}
#[derive(Debug, Serialize, ToSchema)]
pub struct ProviderSecretsResponse {
pub secrets: Vec<ProviderSecret>,
}
#[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
pub enum CommandType {
Builtin,
@@ -248,6 +290,343 @@ fn mask_secret(secret: Value) -> String {
format!("{}{}", visible, mask)
}
const SECRET_STORE_ID_PREFIX: &str = "secret_store:";
const PROVIDER_CACHE_ID_PREFIX: &str = "provider_cache:";
fn provider_secret_status(expires_at: Option<DateTime<Utc>>) -> ProviderSecretStatus {
match expires_at {
Some(expires_at) if expires_at <= Utc::now() => ProviderSecretStatus::Expired,
Some(_) => ProviderSecretStatus::Valid,
None => ProviderSecretStatus::Unknown,
}
}
fn parse_expiry_value(value: &Value) -> Option<DateTime<Utc>> {
match value {
Value::String(value) => DateTime::parse_from_rfc3339(value)
.ok()
.map(|dt| dt.with_timezone(&Utc)),
Value::Number(value) => value
.as_i64()
.and_then(|timestamp| Utc.timestamp_opt(timestamp, 0).single()),
_ => None,
}
}
fn find_expires_at(value: &Value) -> Option<DateTime<Utc>> {
match value {
Value::Object(map) => {
if map
.get("refresh_token")
.and_then(Value::as_str)
.is_some_and(|token| !token.is_empty())
{
return None;
}
if let Some(expires_at) = map.get("expires_at").and_then(parse_expiry_value) {
return Some(expires_at);
}
if let Some(expires_at) = map.get("expires_on").and_then(parse_expiry_value) {
return Some(expires_at);
}
map.values().find_map(find_expires_at)
}
Value::Array(values) => values.iter().find_map(find_expires_at),
_ => None,
}
}
#[derive(Clone, Copy)]
struct ProviderCacheSecretDefinition {
provider: &'static str,
name: &'static str,
path: &'static str,
is_directory: bool,
}
const PROVIDER_CACHE_SECRET_DEFINITIONS: &[ProviderCacheSecretDefinition] = &[
ProviderCacheSecretDefinition {
provider: "gemini_oauth",
name: "OAuth token",
path: "gemini_oauth/tokens.json",
is_directory: false,
},
ProviderCacheSecretDefinition {
provider: "chatgpt_codex",
name: "OAuth token",
path: "chatgpt_codex/tokens.json",
is_directory: false,
},
ProviderCacheSecretDefinition {
provider: "kimi_code",
name: "OAuth token",
path: "kimicode/token.json",
is_directory: false,
},
ProviderCacheSecretDefinition {
provider: "github_copilot",
name: "OAuth token",
path: "githubcopilot",
is_directory: true,
},
ProviderCacheSecretDefinition {
provider: "xai_oauth",
name: "OAuth token",
path: "xai_oauth/tokens.json",
is_directory: false,
},
ProviderCacheSecretDefinition {
provider: "databricks",
name: "OAuth token",
path: "databricks/oauth",
is_directory: true,
},
ProviderCacheSecretDefinition {
provider: "databricks_v2",
name: "OAuth token",
path: "databricks/oauth",
is_directory: true,
},
];
fn provider_cache_definitions_for_display() -> Vec<ProviderCacheSecretDefinition> {
let mut seen_paths = HashSet::new();
PROVIDER_CACHE_SECRET_DEFINITIONS
.iter()
.copied()
.filter(|definition| seen_paths.insert(definition.path))
.collect()
}
fn provider_cache_definition(provider: &str) -> Option<ProviderCacheSecretDefinition> {
PROVIDER_CACHE_SECRET_DEFINITIONS
.iter()
.copied()
.find(|definition| definition.provider == provider)
}
fn provider_cache_providers_sharing_cache(provider: &str) -> Vec<&'static str> {
let Some(definition) = provider_cache_definition(provider) else {
return Vec::new();
};
PROVIDER_CACHE_SECRET_DEFINITIONS
.iter()
.filter(|other| other.path == definition.path)
.map(|definition| definition.provider)
.collect()
}
fn read_json_file(path: &std::path::Path) -> Option<Value> {
std::fs::read_to_string(path)
.ok()
.and_then(|contents| serde_json::from_str(&contents).ok())
}
fn collect_json_expiries(path: &std::path::Path, is_directory: bool) -> Vec<DateTime<Utc>> {
if !is_directory {
return read_json_file(path)
.and_then(|value| find_expires_at(&value))
.into_iter()
.collect();
}
let mut expiries = Vec::new();
let mut stack = vec![path.to_path_buf()];
while let Some(current) = stack.pop() {
let Ok(entries) = std::fs::read_dir(current) else {
continue;
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
stack.push(path);
continue;
}
if path.extension().and_then(|ext| ext.to_str()) != Some("json") {
continue;
}
if let Some(expires_at) =
read_json_file(&path).and_then(|value| find_expires_at(&value))
{
expiries.push(expires_at);
}
}
}
expiries
}
fn provider_cache_exists(path: &std::path::Path, is_directory: bool) -> bool {
if !is_directory {
return path.is_file();
}
let Ok(entries) = std::fs::read_dir(path) else {
return false;
};
entries.flatten().any(|entry| {
let path = entry.path();
path.is_file() || provider_cache_exists(&path, true)
})
}
fn provider_cache_expiry(definition: ProviderCacheSecretDefinition) -> Option<DateTime<Utc>> {
let path = Paths::in_config_dir(definition.path);
let expiries = collect_json_expiries(&path, definition.is_directory);
expiries.into_iter().min()
}
fn build_provider_cache_secret(
definition: ProviderCacheSecretDefinition,
display_names: &HashMap<String, String>,
) -> Option<ProviderSecret> {
let path = Paths::in_config_dir(definition.path);
if !provider_cache_exists(&path, definition.is_directory) {
return None;
}
let expires_at = provider_cache_expiry(definition);
Some(ProviderSecret {
id: format!("{}{}", PROVIDER_CACHE_ID_PREFIX, definition.provider),
provider: definition.provider.to_string(),
provider_display_name: display_names
.get(definition.provider)
.cloned()
.unwrap_or_else(|| definition.provider.to_string()),
name: definition.name.to_string(),
storage: ProviderSecretStorage::ProviderCache,
expires_at,
status: provider_secret_status(expires_at),
configured: true,
has_secret: true,
can_delete: true,
can_configure: false,
configure_provider: None,
})
}
fn build_huggingface_oauth_secret(
token: Option<huggingface_auth::HuggingFaceTokenData>,
) -> ProviderSecret {
let expires_at = token.as_ref().and_then(|token| token.expires_at);
let has_secret = token.is_some();
ProviderSecret {
id: format!(
"{}{}",
PROVIDER_CACHE_ID_PREFIX,
huggingface_auth::HUGGINGFACE_PROVIDER_NAME
),
provider: huggingface_auth::HUGGINGFACE_PROVIDER_NAME.to_string(),
provider_display_name: huggingface_auth::HUGGINGFACE_DISPLAY_NAME.to_string(),
name: huggingface_auth::HUGGINGFACE_OAUTH_TOKEN_NAME.to_string(),
storage: ProviderSecretStorage::ProviderCache,
expires_at,
status: provider_secret_status(expires_at),
configured: has_secret,
has_secret,
can_delete: has_secret,
can_configure: true,
configure_provider: Some(huggingface_auth::HUGGINGFACE_PROVIDER_NAME.to_string()),
}
}
fn build_secret_store_secrets(
stored_secrets: &HashMap<String, Value>,
providers: &[(ProviderMetadata, ProviderType)],
) -> Vec<ProviderSecret> {
let mut secrets = Vec::new();
for (metadata, _) in providers {
for config_key in metadata.config_keys.iter().filter(|key| key.secret) {
if !stored_secrets.contains_key(&config_key.name) {
continue;
}
secrets.push(ProviderSecret {
id: format!(
"{}{}:{}",
SECRET_STORE_ID_PREFIX, metadata.name, config_key.name
),
provider: metadata.name.clone(),
provider_display_name: metadata.display_name.clone(),
name: config_key.name.clone(),
storage: ProviderSecretStorage::SecretStore,
expires_at: None,
status: ProviderSecretStatus::Unknown,
configured: true,
has_secret: true,
can_delete: true,
can_configure: false,
configure_provider: None,
});
}
}
secrets
}
fn is_known_provider_secret(
providers: &[(ProviderMetadata, ProviderType)],
provider: &str,
key: &str,
) -> bool {
providers
.iter()
.filter(|(metadata, _)| metadata.name == provider)
.flat_map(|(metadata, _)| metadata.config_keys.iter())
.any(|config_key| config_key.secret && config_key.name == key)
}
fn unconfigure_provider(config: &Config, provider_name: &str) -> Result<(), ConfigError> {
if let Some(mut entry) = goose::config::get_provider_entry(config, provider_name) {
entry.configured = false;
goose::config::set_provider_entry(config, provider_name, &entry)?;
}
let configured_marker = format!("{}_configured", provider_name);
config.delete(&configured_marker)?;
Ok(())
}
fn mark_provider_configured(config: &Config, provider_name: &str) -> Result<(), ConfigError> {
if let Some(mut entry) = goose::config::get_provider_entry(config, provider_name) {
entry.configured = true;
goose::config::set_provider_entry(config, provider_name, &entry)?;
} else {
let model = if goose::config::get_active_provider(config).as_deref() == Some(provider_name)
{
config.get_goose_model().unwrap_or_default()
} else {
String::new()
};
goose::config::set_provider_entry(
config,
provider_name,
&goose::config::ProviderEntry {
enabled: true,
model,
configured: true,
},
)?;
}
Ok(())
}
fn parse_secret_store_id(id: &str) -> Option<(&str, &str)> {
let rest = id.strip_prefix(SECRET_STORE_ID_PREFIX)?;
let (provider, key) = rest.split_once(':')?;
Some((provider, key))
}
fn parse_provider_cache_id(id: &str) -> Option<&str> {
id.strip_prefix(PROVIDER_CACHE_ID_PREFIX)
}
fn is_valid_provider_name(provider_name: &str) -> bool {
!provider_name.is_empty()
&& provider_name
@@ -255,6 +634,123 @@ fn is_valid_provider_name(provider_name: &str) -> bool {
.all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_')
}
fn should_unconfigure_after_secret_delete(
provider: &str,
key: &str,
has_usable_huggingface_oauth_token: impl FnOnce() -> bool,
) -> bool {
provider == huggingface_auth::HUGGINGFACE_PROVIDER_NAME
&& key == huggingface_auth::HUGGINGFACE_TOKEN_SECRET_KEY
&& !has_usable_huggingface_oauth_token()
}
#[utoipa::path(
get,
path = "/config/provider-secrets",
responses(
(status = 200, description = "Provider secrets retrieved successfully", body = ProviderSecretsResponse),
(status = 500, description = "Internal server error")
)
)]
pub async fn list_provider_secrets() -> Result<Json<ProviderSecretsResponse>, ErrorResponse> {
let config = Config::global();
let stored_secrets = config.all_secrets()?;
let providers = get_providers().await;
let display_names: HashMap<String, String> = providers
.iter()
.map(|(metadata, _)| (metadata.name.clone(), metadata.display_name.clone()))
.collect();
let mut secrets = build_secret_store_secrets(&stored_secrets, &providers);
for definition in provider_cache_definitions_for_display() {
if let Some(secret) = build_provider_cache_secret(definition, &display_names) {
if !secrets.iter().any(|existing| existing.id == secret.id) {
secrets.push(secret);
}
}
}
let huggingface_secret = build_huggingface_oauth_secret(huggingface_auth::load_oauth_token());
if let Some(existing) = secrets
.iter_mut()
.find(|existing| existing.id == huggingface_secret.id)
{
*existing = huggingface_secret;
} else {
secrets.push(huggingface_secret);
}
secrets.sort_by(|a, b| {
a.provider_display_name
.cmp(&b.provider_display_name)
.then_with(|| a.name.cmp(&b.name))
});
Ok(Json(ProviderSecretsResponse { secrets }))
}
#[utoipa::path(
delete,
path = "/config/provider-secrets/{id}",
params(
("id" = String, Path, description = "Provider secret identifier")
),
responses(
(status = 200, description = "Provider secret deleted successfully", body = String),
(status = 400, description = "Invalid provider secret identifier"),
(status = 500, description = "Internal server error")
)
)]
pub async fn delete_provider_secret(Path(id): Path<String>) -> Result<Json<String>, ErrorResponse> {
let config = Config::global();
if let Some((provider, key)) = parse_secret_store_id(&id) {
let providers = get_providers().await;
if !is_known_provider_secret(&providers, provider, key) {
return Err(ErrorResponse::bad_request(format!(
"Invalid provider secret id: '{}'",
id
)));
}
config.delete_secret(key)?;
if should_unconfigure_after_secret_delete(provider, key, || {
huggingface_auth::has_configured_token().unwrap_or(false)
}) {
unconfigure_provider(config, provider)?;
}
return Ok(Json(format!("Deleted provider secret {}", id)));
}
if let Some(provider) = parse_provider_cache_id(&id) {
if provider == huggingface_auth::HUGGINGFACE_PROVIDER_NAME {
huggingface_auth::clear_oauth_token()?;
unconfigure_provider(config, provider)?;
return Ok(Json(format!("Deleted provider secret {}", id)));
}
let cache_definition = provider_cache_definition(provider);
if !is_valid_provider_name(provider) || cache_definition.is_none() {
return Err(ErrorResponse::bad_request(format!(
"Invalid provider name: '{}'",
provider
)));
}
goose::providers::cleanup_provider(provider).await?;
for shared_provider in provider_cache_providers_sharing_cache(provider) {
unconfigure_provider(config, shared_provider)?;
}
return Ok(Json(format!("Deleted provider secret {}", id)));
}
Err(ErrorResponse::bad_request(format!(
"Invalid provider secret id: '{}'",
id
)))
}
#[utoipa::path(
post,
path = "/config/read",
@@ -933,6 +1429,17 @@ pub async fn configure_provider_oauth(
)));
}
if provider_name == huggingface_auth::HUGGINGFACE_PROVIDER_NAME {
huggingface_auth::configure_oauth().await.map_err(|e| {
ErrorResponse::bad_request(format!(
"OAuth configuration failed for provider '{}': {}",
provider_name, e
))
})?;
mark_provider_configured(goose::config::Config::global(), &provider_name)?;
return Ok(Json("OAuth configuration completed".to_string()));
}
let temp_model = ModelConfig::new("temp")
.map_err(|e| {
ErrorResponse::bad_request(format!("Failed to create temporary model config: {}", e))
@@ -956,29 +1463,7 @@ pub async fn configure_provider_oauth(
))
})?;
// Mark the provider as configured after successful OAuth
let config = goose::config::Config::global();
if let Some(mut entry) = goose::config::get_provider_entry(config, &provider_name) {
entry.configured = true;
goose::config::set_provider_entry(config, &provider_name, &entry)?;
} else {
let model = if goose::config::get_active_provider(config).as_deref()
== Some(provider_name.as_str())
{
config.get_goose_model().unwrap_or_default()
} else {
String::new()
};
goose::config::set_provider_entry(
config,
&provider_name,
&goose::config::ProviderEntry {
enabled: true,
model,
configured: true,
},
)?;
}
mark_provider_configured(goose::config::Config::global(), &provider_name)?;
Ok(Json("OAuth configuration completed".to_string()))
}
@@ -989,6 +1474,11 @@ pub fn routes(state: Arc<AppState>) -> Router {
.route("/config/upsert", post(upsert_config))
.route("/config/remove", post(remove_config))
.route("/config/read", post(read_config))
.route("/config/provider-secrets", get(list_provider_secrets))
.route(
"/config/provider-secrets/{id}",
delete(delete_provider_secret),
)
.route("/config/extensions", get(get_extensions))
.route("/config/extensions", post(add_extension))
.route("/config/extensions/{name}", delete(remove_extension))
@@ -1031,4 +1521,281 @@ pub fn routes(state: Arc<AppState>) -> Router {
}
#[cfg(test)]
mod tests {}
mod tests {
use super::*;
use goose::config::ProviderEntry;
use goose::providers::base::ConfigKey;
use serde_json::json;
fn new_test_config() -> Config {
let unique = format!(
"goose-server-config-test-{}-{}",
std::process::id(),
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_nanos()
);
let config_path = std::env::temp_dir().join(format!("{unique}-config.yaml"));
let secrets_path = std::env::temp_dir().join(format!("{unique}-secrets.yaml"));
Config::new_with_file_secrets(config_path, secrets_path).unwrap()
}
#[test]
fn secret_store_listing_only_includes_provider_secret_keys() {
let metadata = ProviderMetadata::new(
"openai",
"OpenAI",
"OpenAI provider",
"gpt-4o",
vec![],
"https://example.com",
vec![
ConfigKey::new("OPENAI_API_KEY", true, true, None, true),
ConfigKey::new("OPENAI_HOST", false, false, None, false),
],
);
let providers = vec![(metadata, ProviderType::Builtin)];
let stored_secrets = HashMap::from([
(
"OPENAI_API_KEY".to_string(),
Value::String("secret-value".to_string()),
),
(
"UNRELATED_SECRET".to_string(),
Value::String("other-secret".to_string()),
),
(
"OPENAI_HOST".to_string(),
Value::String("https://api.openai.com".to_string()),
),
]);
let secrets = build_secret_store_secrets(&stored_secrets, &providers);
assert_eq!(secrets.len(), 1);
assert_eq!(secrets[0].id, "secret_store:openai:OPENAI_API_KEY");
assert_eq!(secrets[0].provider_display_name, "OpenAI");
assert_eq!(secrets[0].name, "OPENAI_API_KEY");
assert_eq!(secrets[0].storage, ProviderSecretStorage::SecretStore);
assert_eq!(secrets[0].status, ProviderSecretStatus::Unknown);
}
#[test]
fn provider_secret_delete_validation_requires_provider_secret_key() {
let metadata = ProviderMetadata::new(
"openai",
"OpenAI",
"OpenAI provider",
"gpt-4o",
vec![],
"https://example.com",
vec![
ConfigKey::new("OPENAI_API_KEY", true, true, None, true),
ConfigKey::new("OPENAI_HOST", false, false, None, false),
],
);
let providers = vec![(metadata, ProviderType::Builtin)];
assert!(is_known_provider_secret(
&providers,
"openai",
"OPENAI_API_KEY"
));
assert!(!is_known_provider_secret(
&providers,
"openai",
"OPENAI_HOST"
));
assert!(!is_known_provider_secret(
&providers,
"openai",
"UNRELATED_SECRET"
));
assert!(!is_known_provider_secret(
&providers,
"anthropic",
"OPENAI_API_KEY"
));
}
#[test]
fn expiry_extraction_handles_nested_rfc3339_values() {
let expires_at = Utc::now() + chrono::Duration::hours(1);
let value = json!({
"project_id": "project",
"token": {
"access_token": "secret",
"expires_at": expires_at.to_rfc3339(),
}
});
let parsed = find_expires_at(&value).expect("expected expiry");
assert_eq!(parsed.timestamp(), expires_at.timestamp());
assert_eq!(
provider_secret_status(Some(parsed)),
ProviderSecretStatus::Valid
);
}
#[test]
fn expiry_extraction_ignores_refreshable_access_tokens() {
let expires_at = Utc::now() - chrono::Duration::hours(1);
let value = json!({
"access_token": "access",
"refresh_token": "refresh",
"expires_at": expires_at.to_rfc3339(),
});
assert_eq!(find_expires_at(&value), None);
}
#[test]
fn expiry_extraction_handles_expired_unix_timestamps() {
let value = json!({
"info": {
"expires_at": 1
}
});
let parsed = find_expires_at(&value).expect("expected expiry");
assert_eq!(parsed.timestamp(), 1);
assert_eq!(
provider_secret_status(Some(parsed)),
ProviderSecretStatus::Expired
);
}
#[test]
fn provider_secret_ids_parse_expected_prefixes() {
assert_eq!(
parse_secret_store_id("secret_store:openai:OPENAI_API_KEY"),
Some(("openai", "OPENAI_API_KEY"))
);
assert_eq!(
parse_provider_cache_id("provider_cache:gemini_oauth"),
Some("gemini_oauth")
);
assert_eq!(parse_secret_store_id("provider_cache:openai"), None);
assert_eq!(parse_provider_cache_id("secret_store:openai:key"), None);
}
#[test]
fn shared_databricks_cache_is_displayed_once() {
let databricks_definitions: Vec<_> = provider_cache_definitions_for_display()
.into_iter()
.filter(|definition| definition.path == "databricks/oauth")
.collect();
assert_eq!(databricks_definitions.len(), 1);
assert_eq!(databricks_definitions[0].provider, "databricks");
}
#[test]
fn shared_databricks_cache_unconfigures_both_providers() {
assert_eq!(
provider_cache_providers_sharing_cache("databricks"),
vec!["databricks", "databricks_v2"]
);
assert_eq!(
provider_cache_providers_sharing_cache("databricks_v2"),
vec!["databricks", "databricks_v2"]
);
}
#[test]
fn unconfigure_provider_clears_structured_entry() {
let config = new_test_config();
goose::config::set_provider_entry(
&config,
"huggingface",
&ProviderEntry {
enabled: true,
model: "Qwen/Qwen3-Coder-480B-A35B-Instruct".to_string(),
configured: true,
},
)
.unwrap();
unconfigure_provider(&config, "huggingface").unwrap();
let entry = goose::config::get_provider_entry(&config, "huggingface").unwrap();
assert!(entry.enabled);
assert_eq!(entry.model, "Qwen/Qwen3-Coder-480B-A35B-Instruct");
assert!(!entry.configured);
}
#[test]
fn unconfigure_provider_deletes_legacy_configured_marker() {
let config = new_test_config();
config.set_param("huggingface_configured", true).unwrap();
unconfigure_provider(&config, "huggingface").unwrap();
assert!(config.get_param::<bool>("huggingface_configured").is_err());
}
#[test]
fn deleting_huggingface_token_unconfigures_without_oauth() {
assert!(should_unconfigure_after_secret_delete(
"huggingface",
"HF_TOKEN",
|| false
));
}
#[test]
fn deleting_huggingface_token_keeps_configured_with_oauth() {
assert!(!should_unconfigure_after_secret_delete(
"huggingface",
"HF_TOKEN",
|| true
));
}
#[test]
fn deleting_other_provider_secret_does_not_unconfigure_huggingface() {
assert!(!should_unconfigure_after_secret_delete(
"openai",
"OPENAI_API_KEY",
|| false
));
}
#[test]
fn huggingface_oauth_secret_is_permanent_without_token() {
let secret = build_huggingface_oauth_secret(None);
assert_eq!(secret.id, "provider_cache:huggingface");
assert_eq!(secret.provider_display_name, "Hugging Face");
assert_eq!(secret.name, "OAuth token");
assert_eq!(secret.storage, ProviderSecretStorage::ProviderCache);
assert_eq!(secret.status, ProviderSecretStatus::Unknown);
assert!(!secret.configured);
assert!(!secret.has_secret);
assert!(!secret.can_delete);
assert!(secret.can_configure);
assert_eq!(secret.configure_provider.as_deref(), Some("huggingface"));
}
#[test]
fn huggingface_oauth_secret_reports_cached_token_metadata() {
let expires_at = Utc::now() + chrono::Duration::hours(1);
let secret = build_huggingface_oauth_secret(Some(huggingface_auth::HuggingFaceTokenData {
access_token: "hidden".to_string(),
refresh_token: None,
expires_at: Some(expires_at),
}));
assert_eq!(
secret.expires_at.map(|value| value.timestamp()),
Some(expires_at.timestamp())
);
assert_eq!(secret.status, ProviderSecretStatus::Valid);
assert!(secret.configured);
assert!(secret.has_secret);
assert!(secret.can_delete);
}
}
crates/goose-server/src/routes/local_inference.rs
+17 -3
View File
@@ -11,6 +11,7 @@ use axum::{
use futures::future::join_all;
use goose::config::paths::Paths;
use goose::download_manager::{get_download_manager, DownloadProgress};
use goose::providers::huggingface_auth;
use goose::providers::local_inference::hf_models::{self, HfModelInfo, HfQuantVariant};
use goose::providers::local_inference::{
available_inference_memory_bytes, builtin_chat_template_names,
@@ -251,6 +252,7 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
// Auto-download mmproj files for models that are already downloaded.
// Deduplicate by path since multiple quants share one mmproj file.
let dm = get_download_manager();
let hf_token = huggingface_auth::resolve_token_async().await.ok().flatten();
let mut started_paths = std::collections::HashSet::new();
for (model_id, url, path) in mmproj_downloads_needed {
if !path.exists() && started_paths.insert(path.clone()) {
@@ -260,7 +262,16 @@ async fn ensure_featured_models_in_registry() -> Result<(), ErrorResponse> {
.is_some_and(|p| p.status == goose::download_manager::DownloadStatus::Downloading);
if !dominated_by_active {
tracing::info!(model_id = %model_id, "Auto-downloading vision encoder for existing model");
if let Err(e) = dm.download_model(download_id, url, path, None).await {
if let Err(e) = dm
.download_model_with_bearer_token(
download_id,
url,
path,
hf_token.clone(),
None,
)
.await
{
tracing::warn!(model_id = %model_id, error = %e, "Failed to start mmproj download");
}
}
@@ -471,6 +482,7 @@ pub async fn download_hf_model(
let (_repo, resolved) = resolve_model_spec_full(&req.spec)
.await
.map_err(|e| ErrorResponse::bad_request(format!("Invalid spec: {}", e)))?;
let hf_token = huggingface_auth::resolve_token_async().await.ok().flatten();
let model_id = model_id_from_repo(&repo_id, &quantization);
let models_dir = Paths::in_data_dir("models");
@@ -545,10 +557,11 @@ pub async fn download_hf_model(
.map(|f| (f.download_url.clone(), models_dir.join(&f.filename)))
.collect();
dm.download_model_sharded(
dm.download_model_sharded_with_bearer_token(
format!("{}-model", model_id),
all_files,
resolved.total_size,
hf_token.clone(),
None,
)
.await
@@ -556,10 +569,11 @@ pub async fn download_hf_model(
if let Some((mmproj_path, mmproj_url)) = mmproj_path {
if !mmproj_path.exists() {
dm.download_model(
dm.download_model_with_bearer_token(
format!("{}-mmproj", model_id),
mmproj_url,
mmproj_path,
hf_token,
None,
)
.await
crates/goose-server/src/routes/utils.rs
+111 -2
View File
@@ -1,6 +1,7 @@
use goose::config::declarative_providers::load_provider;
use goose::config::declarative_providers::{load_provider, LoadedProvider};
use goose::config::Config;
use goose::providers::base::{ConfigKey, ProviderMetadata, ProviderType};
use goose::providers::huggingface_auth;
use serde::{Deserialize, Serialize};
use std::env;
use std::error::Error;
@@ -92,15 +93,37 @@ pub fn inspect_keys(
}
pub fn check_provider_configured(metadata: &ProviderMetadata, provider_type: ProviderType) -> bool {
let config = Config::global();
check_provider_configured_with_huggingface_oauth(metadata, provider_type, || {
huggingface_auth::has_usable_or_refreshable_oauth_token()
})
}
fn check_provider_configured_with_huggingface_oauth(
metadata: &ProviderMetadata,
provider_type: ProviderType,
has_usable_huggingface_oauth_token: impl Fn() -> bool,
) -> bool {
// Special override
if metadata.name == "local" {
return true;
}
if accepts_huggingface_oauth(metadata, None, &has_usable_huggingface_oauth_token) {
return true;
}
let config = Config::global();
if provider_type == ProviderType::Custom || provider_type == ProviderType::Declarative {
if let Ok(loaded_provider) = load_provider(metadata.name.as_str()) {
if accepts_huggingface_oauth(
metadata,
Some(&loaded_provider),
&has_usable_huggingface_oauth_token,
) {
return true;
}
if !loaded_provider.config.requires_auth {
return true;
}
@@ -211,3 +234,89 @@ pub fn check_provider_configured(metadata: &ProviderMetadata, provider_type: Pro
is_set_in_env || is_set_in_config
})
}
fn accepts_huggingface_oauth(
metadata: &ProviderMetadata,
loaded_provider: Option<&LoadedProvider>,
has_usable_huggingface_oauth_token: &impl Fn() -> bool,
) -> bool {
let is_huggingface_provider = metadata.name == huggingface_auth::HUGGINGFACE_PROVIDER_NAME
|| loaded_provider.is_some_and(|provider| {
provider.config.catalog_provider_id.as_deref()
== Some(huggingface_auth::HUGGINGFACE_PROVIDER_NAME)
});
is_huggingface_provider && has_usable_huggingface_oauth_token()
}
#[cfg(test)]
mod tests {
use super::*;
use goose::config::declarative_providers::{DeclarativeProviderConfig, ProviderEngine};
use goose::providers::base::ModelInfo;
fn huggingface_metadata() -> ProviderMetadata {
ProviderMetadata::new(
huggingface_auth::HUGGINGFACE_PROVIDER_NAME,
huggingface_auth::HUGGINGFACE_DISPLAY_NAME,
"Hugging Face provider",
"Qwen/Qwen3-Coder-480B-A35B-Instruct",
vec![],
"https://huggingface.co/docs/inference-providers",
vec![ConfigKey::new(
huggingface_auth::HUGGINGFACE_TOKEN_SECRET_KEY,
true,
true,
None,
true,
)],
)
}
#[test]
fn huggingface_oauth_token_counts_as_configured_without_hf_token() {
assert!(check_provider_configured_with_huggingface_oauth(
&huggingface_metadata(),
ProviderType::Builtin,
|| true,
));
}
#[test]
fn huggingface_catalog_provider_oauth_counts_as_configured() {
let mut metadata = huggingface_metadata();
metadata.name = "custom-huggingface".to_string();
let loaded_provider = LoadedProvider {
config: DeclarativeProviderConfig {
name: metadata.name.clone(),
engine: ProviderEngine::OpenAI,
display_name: "Custom Hugging Face".to_string(),
description: None,
api_key_env: String::new(),
base_url: "https://router.huggingface.co/v1".to_string(),
models: vec![ModelInfo::new("test-model", 128_000)],
headers: None,
timeout_seconds: None,
supports_streaming: None,
requires_auth: true,
catalog_provider_id: Some(huggingface_auth::HUGGINGFACE_PROVIDER_NAME.to_string()),
base_path: None,
env_vars: None,
dynamic_models: None,
skip_canonical_filtering: false,
model_doc_link: None,
setup_steps: vec![],
fast_model: None,
preserves_thinking: false,
},
is_editable: false,
};
assert!(accepts_huggingface_oauth(
&metadata,
Some(&loaded_provider),
&|| true,
));
}
}
crates/goose/src/config/declarative_providers.rs
+154 -16
View File
@@ -1,7 +1,8 @@
use crate::config::paths::Paths;
use crate::config::Config;
use crate::providers::anthropic::AnthropicProvider;
use crate::providers::base::{ModelInfo, ProviderType};
use crate::providers::base::{ModelInfo, ProviderDef, ProviderType};
use crate::providers::huggingface::HuggingFaceProvider;
use crate::providers::inventory::declarative_inventory_identity;
use crate::providers::ollama::OllamaProvider;
use crate::providers::openai::OpenAiProvider;
@@ -576,21 +577,48 @@ pub fn register_declarative_provider(
ProviderEngine::OpenAI => {
let captured = config.clone();
let identity_config = config.clone();
registry.register_with_name::<OpenAiProvider, _, _>(
&config,
provider_type,
config.dynamic_models.unwrap_or(false),
move |model| {
let mut cfg = captured.clone();
resolve_config(&mut cfg)?;
OpenAiProvider::from_custom_config(model, cfg)
},
move || {
let mut cfg = identity_config.clone();
resolve_config(&mut cfg)?;
declarative_inventory_identity(&cfg)
},
);
if HuggingFaceProvider::matches_declarative_config(&config) {
let inventory_configured_config = config.clone();
registry
.register_with_name_and_inventory_configured::<HuggingFaceProvider, _, _, _>(
&config,
provider_type,
config.dynamic_models.unwrap_or(false),
move |model| {
let mut cfg = captured.clone();
resolve_config(&mut cfg)?;
HuggingFaceProvider::from_custom_config(model, cfg)
},
move || {
let mut cfg = identity_config.clone();
resolve_config(&mut cfg)?;
declarative_inventory_identity(&cfg)
},
move || {
let mut cfg = inventory_configured_config.clone();
if resolve_config(&mut cfg).is_err() {
return false;
}
huggingface_declarative_inventory_configured(&cfg)
},
);
} else {
registry.register_with_name::<OpenAiProvider, _, _>(
&config,
provider_type,
config.dynamic_models.unwrap_or(false),
move |model| {
let mut cfg = captured.clone();
resolve_config(&mut cfg)?;
OpenAiProvider::from_custom_config(model, cfg)
},
move || {
let mut cfg = identity_config.clone();
resolve_config(&mut cfg)?;
declarative_inventory_identity(&cfg)
},
);
}
}
ProviderEngine::Ollama => {
let captured = config.clone();
@@ -633,10 +661,120 @@ pub fn register_declarative_provider(
}
}
fn huggingface_declarative_inventory_configured(config: &DeclarativeProviderConfig) -> bool {
huggingface_declarative_inventory_configured_from_sources(
config,
|key| Config::global().get_secret::<String>(key).is_ok(),
HuggingFaceProvider::inventory_configured,
)
}
fn huggingface_declarative_inventory_configured_from_sources(
config: &DeclarativeProviderConfig,
provider_secret_configured: impl FnOnce(&str) -> bool,
global_huggingface_configured: impl FnOnce() -> bool,
) -> bool {
if !config.requires_auth {
return true;
}
if !config.api_key_env.is_empty() {
return provider_secret_configured(&config.api_key_env);
}
global_huggingface_configured()
}
#[cfg(test)]
mod tests {
use super::*;
fn test_huggingface_config() -> DeclarativeProviderConfig {
DeclarativeProviderConfig {
name: "custom_hf".to_string(),
engine: ProviderEngine::OpenAI,
display_name: "Custom HF".to_string(),
description: None,
api_key_env: String::new(),
base_url: "https://router.huggingface.co/v1".to_string(),
models: vec![ModelInfo {
name: "test/model".to_string(),
resolved_model: None,
context_limit: 128_000,
input_token_cost: None,
output_token_cost: None,
currency: None,
supports_cache_control: None,
reasoning: false,
}],
headers: None,
timeout_seconds: None,
supports_streaming: Some(true),
requires_auth: true,
catalog_provider_id: Some("huggingface".to_string()),
base_path: None,
env_vars: None,
dynamic_models: Some(false),
skip_canonical_filtering: false,
model_doc_link: None,
setup_steps: Vec::new(),
fast_model: None,
preserves_thinking: true,
}
}
#[test]
fn huggingface_inventory_allows_unauthenticated_custom_provider() {
let mut config = test_huggingface_config();
config.requires_auth = false;
assert!(huggingface_declarative_inventory_configured_from_sources(
&config,
|_| false,
|| false,
));
}
#[test]
fn huggingface_inventory_accepts_provider_specific_key() {
let mut config = test_huggingface_config();
config.api_key_env = "CUSTOM_HF_TOKEN".to_string();
assert!(huggingface_declarative_inventory_configured_from_sources(
&config,
|key| key == "CUSTOM_HF_TOKEN",
|| false,
));
}
#[test]
fn huggingface_inventory_does_not_fallback_when_explicit_key_is_missing() {
let mut config = test_huggingface_config();
config.api_key_env = "CUSTOM_HF_TOKEN".to_string();
assert!(!huggingface_declarative_inventory_configured_from_sources(
&config,
|_| false,
|| true,
));
}
#[test]
fn huggingface_inventory_uses_global_token_without_provider_key() {
let config = test_huggingface_config();
assert!(huggingface_declarative_inventory_configured_from_sources(
&config,
|_| false,
|| true,
));
assert!(!huggingface_declarative_inventory_configured_from_sources(
&config,
|_| true,
|| false,
));
}
#[test]
fn test_tanzu_json_deserializes() {
let json = include_str!("../providers/declarative/tanzu.json");
crates/goose/src/download_manager.rs
+60 -7
View File
@@ -131,12 +131,48 @@ impl DownloadManager {
.await
}
pub async fn download_model_with_bearer_token(
&self,
model_id: String,
url: String,
destination: PathBuf,
bearer_token: Option<String>,
on_complete: Option<Box<dyn FnOnce() + Send + 'static>>,
) -> Result<()> {
self.download_model_sharded_with_bearer_token(
model_id,
vec![(url, destination)],
0,
bearer_token,
on_complete,
)
.await
}
pub async fn download_model_sharded(
&self,
model_id: String,
files: Vec<(String, PathBuf)>,
total_size_hint: u64,
on_complete: Option<Box<dyn FnOnce() + Send + 'static>>,
) -> Result<()> {
self.download_model_sharded_with_bearer_token(
model_id,
files,
total_size_hint,
None,
on_complete,
)
.await
}
pub async fn download_model_sharded_with_bearer_token(
&self,
model_id: String,
files: Vec<(String, PathBuf)>,
total_size_hint: u64,
bearer_token: Option<String>,
on_complete: Option<Box<dyn FnOnce() + Send + 'static>>,
) -> Result<()> {
info!(model_id = %model_id, file_count = files.len(), "Starting model download");
{
@@ -186,8 +222,13 @@ impl DownloadManager {
let files_for_cleanup: Vec<PathBuf> = files.iter().map(|(_, d)| d.clone()).collect();
tokio::spawn(async move {
let result =
Self::download_files_sequentially(&files, &downloads, &model_id_clone).await;
let result = Self::download_files_sequentially(
&files,
&downloads,
&model_id_clone,
bearer_token.as_deref(),
)
.await;
match result {
Ok(_) => {
@@ -262,6 +303,7 @@ impl DownloadManager {
files: &[(String, PathBuf)],
downloads: &DownloadMap,
model_id: &str,
bearer_token: Option<&str>,
) -> Result<(), anyhow::Error> {
let client = reqwest::Client::builder()
.connect_timeout(std::time::Duration::from_secs(30))
@@ -273,8 +315,7 @@ impl DownloadManager {
let mut total: u64 = 0;
let mut all_resolved = true;
for (url, _) in files {
let size = client
.head(url)
let size = Self::apply_bearer_token(client.head(url), bearer_token)
.send()
.await
.ok()
@@ -329,6 +370,7 @@ impl DownloadManager {
&mut cumulative_bytes,
start_time,
bytes_at_start,
bearer_token,
)
.await?;
}
@@ -346,6 +388,7 @@ impl DownloadManager {
cumulative_bytes: &mut u64,
start_time: std::time::Instant,
bytes_at_start: u64,
bearer_token: Option<&str>,
) -> Result<(), anyhow::Error> {
let partial_path = partial_path_for(destination);
let mut retries = 0u32;
@@ -357,8 +400,7 @@ impl DownloadManager {
};
// Get this file's total size
let mut file_total: u64 = client
.head(url)
let mut file_total: u64 = Self::apply_bearer_token(client.head(url), bearer_token)
.send()
.await
.ok()
@@ -386,7 +428,7 @@ impl DownloadManager {
anyhow::bail!("Download cancelled");
}
let mut request = client.get(url);
let mut request = Self::apply_bearer_token(client.get(url), bearer_token);
if file_bytes > 0 {
request = request.header("Range", format!("bytes={}-", file_bytes));
}
@@ -584,6 +626,17 @@ impl DownloadManager {
}
}
}
fn apply_bearer_token(
request: reqwest::RequestBuilder,
bearer_token: Option<&str>,
) -> reqwest::RequestBuilder {
if let Some(token) = bearer_token.filter(|token| !token.is_empty()) {
request.header("Authorization", format!("Bearer {}", token))
} else {
request
}
}
}
static DOWNLOAD_MANAGER: once_cell::sync::Lazy<DownloadManager> =
crates/goose/src/providers/catalog.rs
+31
View File
@@ -375,6 +375,23 @@ const SETUP_METADATA: &[CuratedSetupMetadata] = &[
secret_field_default: Some(API_KEY_FIELD),
field_overrides: &[],
},
CuratedSetupMetadata {
provider_id: "huggingface",
category: ProviderSetupCategory::Model,
setup_method: ProviderSetupMethod::SingleApiKey,
group: ProviderSetupGroup::Default,
display_name: Some("Hugging Face"),
description: Some("Hugging Face Inference Providers"),
docs_url: Some("https://huggingface.co/docs/inference-providers"),
aliases: &["huggingface", "hf"],
native_connect_query: None,
binary_name: None,
setup_capabilities: setup_capabilities(false, false, false),
show_only_when_installed: false,
synthetic: false,
secret_field_default: Some(API_KEY_FIELD),
field_overrides: &[],
},
CuratedSetupMetadata {
provider_id: "chatgpt_codex",
category: ProviderSetupCategory::Model,
@@ -1199,6 +1216,20 @@ mod tests {
["DATABRICKS_HOST", "DATABRICKS_TOKEN"]
);
let huggingface = entries
.iter()
.find(|entry| entry.provider_id == "huggingface")
.expect("setup catalog should include huggingface");
assert_eq!(huggingface.setup_method, ProviderSetupMethod::SingleApiKey);
assert_eq!(
huggingface
.fields
.iter()
.map(|field| field.key.as_str())
.collect::<Vec<_>>(),
["HF_TOKEN"]
);
let atomic_chat = entries
.iter()
.find(|entry| entry.provider_id == "atomic_chat")
crates/goose/src/providers/gemini_oauth.rs
+5
View File
@@ -851,6 +851,11 @@ impl GeminiOAuthProvider {
})
}
pub async fn cleanup() -> Result<()> {
TokenCache::new().clear();
Ok(())
}
async fn post_stream(
&self,
session_id: Option<&str>,
crates/goose/src/providers/huggingface.rs
+569
View File
@@ -0,0 +1,569 @@
use super::api_client::{ApiClient, AuthMethod, AuthProvider};
use super::base::{
ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata,
DEFAULT_PROVIDER_TIMEOUT_SECS,
};
use super::errors::ProviderError;
use super::huggingface_auth;
use super::inventory::{default_inventory_identity, InventoryIdentityInput};
use super::openai_compatible::OpenAiCompatibleProvider;
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::config::{Config, ConfigError};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use anyhow::{anyhow, Result};
use futures::future::BoxFuture;
use rmcp::model::Tool;
pub const HUGGINGFACE_API_HOST: &str = "https://router.huggingface.co/v1";
pub const HUGGINGFACE_DOC_URL: &str = "https://huggingface.co/docs/inference-providers";
pub const HUGGINGFACE_DEFAULT_MODEL: &str = "Qwen/Qwen3-Coder-480B-A35B-Instruct";
pub const HUGGINGFACE_KNOWN_MODELS: &[&str] = &[
"MiniMaxAI/MiniMax-M2.1",
"MiniMaxAI/MiniMax-M2.5",
"MiniMaxAI/MiniMax-M2.7",
"Qwen/Qwen3-235B-A22B-Thinking",
"Qwen/Qwen3-Coder-480B-A35B-Instruct",
"Qwen/Qwen3-Coder-Next",
"Qwen/Qwen3-Embedding-4B",
"Qwen/Qwen3-Embedding-8B",
"Qwen/Qwen3-Next-80B-A3B-Instruct",
"Qwen/Qwen3-Next-80B-A3B-Thinking",
"Qwen/Qwen3.5-397B-A17B",
"XiaomiMiMo/MiMo-V2-Flash",
"deepseek-ai/DeepSeek-R1",
"deepseek-ai/DeepSeek-V3.2",
"deepseek-ai/DeepSeek-V4-Pro",
"moonshotai/Kimi-K2-Instruct",
"moonshotai/Kimi-K2-Thinking",
"moonshotai/Kimi-K2.5",
"moonshotai/Kimi-K2.6",
"zai-org/GLM-4.7",
"zai-org/GLM-4.7-Flash",
"zai-org/GLM-5",
"zai-org/GLM-5.1",
];
type QueryParams = Vec<(String, String)>;
type EndpointParts = (String, String, QueryParams);
pub struct HuggingFaceProvider {
inner: OpenAiCompatibleProvider,
custom_models: Option<Vec<String>>,
dynamic_models: Option<bool>,
}
struct HuggingFaceAuthProvider;
#[async_trait::async_trait]
impl AuthProvider for HuggingFaceAuthProvider {
async fn get_auth_header(&self) -> Result<(String, String)> {
let token = huggingface_auth::resolve_token_async()
.await?
.ok_or_else(missing_token_error)?;
Ok(("Authorization".to_string(), format!("Bearer {}", token)))
}
}
impl HuggingFaceProvider {
pub fn matches_declarative_config(config: &DeclarativeProviderConfig) -> bool {
config.name == huggingface_auth::HUGGINGFACE_PROVIDER_NAME
|| config.catalog_provider_id.as_deref()
== Some(huggingface_auth::HUGGINGFACE_PROVIDER_NAME)
}
pub fn from_custom_config(
model: ModelConfig,
config: DeclarativeProviderConfig,
) -> Result<Self> {
let custom_models = static_model_names(&config);
if config.dynamic_models == Some(false) && custom_models.is_none() {
return Err(anyhow!(
"Provider '{}' has dynamic_models: false but no static models listed; \
at least one entry in `models` is required.",
config.name
));
}
let auth_method = custom_auth_method(&config)?;
let (host, completions_prefix, query_params) =
openai_compatible_endpoint_parts(&config.base_url, config.base_path.as_deref())?;
let timeout_secs = config
.timeout_seconds
.unwrap_or(DEFAULT_PROVIDER_TIMEOUT_SECS);
let mut api_client = ApiClient::with_timeout(
host,
auth_method,
std::time::Duration::from_secs(timeout_secs),
)?
.with_query(query_params);
if let Some(headers) = &config.headers {
let mut header_map = reqwest::header::HeaderMap::new();
for (key, value) in headers {
let header_name = reqwest::header::HeaderName::from_bytes(key.as_bytes())?;
let header_value = reqwest::header::HeaderValue::from_str(value)?;
header_map.insert(header_name, header_value);
}
api_client = api_client.with_headers(header_map)?;
}
let model = if let Some(ref fast_model_name) = config.fast_model {
model.with_fast(fast_model_name, &config.name)?
} else {
model
};
Ok(Self {
inner: OpenAiCompatibleProvider::new(
config.name.clone(),
api_client,
model,
completions_prefix,
)
.with_supports_streaming(config.supports_streaming.unwrap_or(true)),
custom_models,
dynamic_models: config.dynamic_models,
})
}
pub async fn cleanup() -> Result<()> {
huggingface_auth::clear_oauth_token()
}
}
#[async_trait::async_trait]
impl Provider for HuggingFaceProvider {
fn get_name(&self) -> &str {
self.inner.get_name()
}
fn get_model_config(&self) -> ModelConfig {
self.inner.get_model_config()
}
async fn fetch_supported_models(&self) -> Result<Vec<String>, ProviderError> {
if let Some(custom_models) = &self.custom_models {
if self.dynamic_models == Some(false) {
return Ok(custom_models.clone());
}
match self.inner.fetch_supported_models().await {
Ok(models) => return Ok(models),
Err(e) if e.is_endpoint_not_found() => {
tracing::debug!(
"Models endpoint not implemented for Hugging Face provider '{}' ({}), using predefined list",
self.inner.get_name(),
e
);
return Ok(custom_models.clone());
}
Err(e) => return Err(e),
}
}
self.inner.fetch_supported_models().await
}
async fn stream(
&self,
model_config: &ModelConfig,
session_id: &str,
system: &str,
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
self.inner
.stream(model_config, session_id, system, messages, tools)
.await
}
}
impl ProviderDef for HuggingFaceProvider {
type Provider = Self;
fn metadata() -> ProviderMetadata {
ProviderMetadata::new(
huggingface_auth::HUGGINGFACE_PROVIDER_NAME,
huggingface_auth::HUGGINGFACE_DISPLAY_NAME,
"Hugging Face Inference Providers via the Hugging Face Router",
HUGGINGFACE_DEFAULT_MODEL,
HUGGINGFACE_KNOWN_MODELS.to_vec(),
HUGGINGFACE_DOC_URL,
vec![
ConfigKey::new(
huggingface_auth::HUGGINGFACE_TOKEN_SECRET_KEY,
true,
true,
None,
true,
),
ConfigKey::new("HF_HOST", false, false, Some(HUGGINGFACE_API_HOST), false),
],
)
}
fn from_env(
model: ModelConfig,
_extensions: Vec<crate::config::ExtensionConfig>,
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(async move {
let config = Config::global();
let auth_method =
refreshable_huggingface_auth_method(huggingface_auth::has_configured_token)?;
let host: String = config
.get_param("HF_HOST")
.unwrap_or_else(|_| HUGGINGFACE_API_HOST.to_string());
let api_client = ApiClient::new(host, auth_method)?;
Ok(Self {
inner: OpenAiCompatibleProvider::new(
huggingface_auth::HUGGINGFACE_PROVIDER_NAME.to_string(),
api_client,
model,
String::new(),
),
custom_models: None,
dynamic_models: None,
})
})
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
let metadata = Self::metadata();
Ok(default_inventory_identity(
&metadata.name,
&metadata.name,
&metadata.config_keys,
Config::global(),
))
}
fn inventory_configured() -> bool {
huggingface_auth::has_configured_token().unwrap_or(false)
}
}
fn missing_token_error() -> anyhow::Error {
anyhow!(
"Hugging Face token is not configured. Sign in from Settings > Auth or configure HF_TOKEN."
)
}
fn configured_api_key(config: &DeclarativeProviderConfig) -> Result<Option<String>> {
if config.api_key_env.is_empty() {
return Ok(None);
}
match Config::global().get_secret::<String>(&config.api_key_env) {
Ok(token) => Ok(Some(token)),
Err(ConfigError::NotFound(_)) => Ok(None),
Err(error) => Err(error.into()),
}
}
fn static_model_names(config: &DeclarativeProviderConfig) -> Option<Vec<String>> {
(!config.models.is_empty()).then(|| {
config
.models
.iter()
.map(|model| model.name.clone())
.collect()
})
}
fn custom_auth_method(config: &DeclarativeProviderConfig) -> Result<AuthMethod> {
let configured_key = if config.requires_auth {
configured_api_key(config)?
} else {
None
};
custom_auth_method_with_provider_token(config.requires_auth, configured_key)
}
fn custom_auth_method_with_provider_token(
requires_auth: bool,
provider_token: Option<String>,
) -> Result<AuthMethod> {
custom_auth_method_from_sources(
requires_auth,
provider_token,
huggingface_auth::has_configured_token,
)
}
fn custom_auth_method_from_sources(
requires_auth: bool,
provider_token: Option<String>,
has_global_token: impl FnOnce() -> Result<bool>,
) -> Result<AuthMethod> {
if !requires_auth {
return Ok(AuthMethod::NoAuth);
}
if let Some(token) = provider_token {
return Ok(AuthMethod::BearerToken(token));
}
refreshable_huggingface_auth_method(has_global_token)
}
fn refreshable_huggingface_auth_method(
has_configured_token: impl FnOnce() -> Result<bool>,
) -> Result<AuthMethod> {
if !has_configured_token()? {
return Err(missing_token_error());
}
Ok(AuthMethod::Custom(Box::new(HuggingFaceAuthProvider)))
}
fn openai_compatible_endpoint_parts(
base_url: &str,
base_path: Option<&str>,
) -> Result<EndpointParts> {
let url =
url::Url::parse(base_url).map_err(|e| anyhow!("Invalid base URL '{}': {}", base_url, e))?;
let mut host = if let Some(port) = url.port() {
format!(
"{}://{}:{}",
url.scheme(),
url.host_str().unwrap_or_default(),
port
)
} else {
format!("{}://{}", url.scheme(), url.host_str().unwrap_or_default())
};
let query_params = url
.query_pairs()
.map(|(key, value)| (key.into_owned(), value.into_owned()))
.collect::<Vec<_>>();
if let Some(path) = base_path {
return Ok((host, completions_prefix(path), query_params));
}
let path = url.path().trim_matches('/');
if path.is_empty() {
return Ok((host, String::new(), query_params));
}
if let Some(parent) = path
.strip_suffix("/chat/completions")
.or_else(|| (path == "chat/completions").then_some(""))
{
if !parent.is_empty() {
host.push('/');
host.push_str(parent);
}
return Ok((host, String::new(), query_params));
}
host.push('/');
host.push_str(path);
Ok((host, String::new(), query_params))
}
fn completions_prefix(path: &str) -> String {
let path = path.trim_matches('/');
if path.is_empty() {
return String::new();
}
let parent = path
.strip_suffix("/chat/completions")
.or_else(|| (path == "chat/completions").then_some(""))
.unwrap_or(path);
if parent.is_empty() {
String::new()
} else {
format!("{}/", parent)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::providers::base::ModelInfo;
#[test]
fn metadata_preserves_huggingface_id_and_token_key() {
let metadata = HuggingFaceProvider::metadata();
assert_eq!(metadata.name, "huggingface");
assert_eq!(metadata.display_name, "Hugging Face");
assert_eq!(metadata.default_model, HUGGINGFACE_DEFAULT_MODEL);
assert!(metadata
.config_keys
.iter()
.any(|key| key.name == "HF_TOKEN" && key.secret));
}
#[test]
fn declarative_matching_accepts_name_or_catalog_provider_id() {
let mut config = test_config();
assert!(!HuggingFaceProvider::matches_declarative_config(&config));
config.name = "huggingface".to_string();
assert!(HuggingFaceProvider::matches_declarative_config(&config));
config.name = "custom_hugging_face".to_string();
config.catalog_provider_id = Some("huggingface".to_string());
assert!(HuggingFaceProvider::matches_declarative_config(&config));
}
#[test]
fn endpoint_parts_use_base_url_path_as_api_host() {
let (host, prefix, query) =
openai_compatible_endpoint_parts("https://router.huggingface.co/v1?beta=1", None)
.unwrap();
assert_eq!(host, "https://router.huggingface.co/v1");
assert_eq!(prefix, "");
assert_eq!(query, vec![("beta".to_string(), "1".to_string())]);
}
#[test]
fn endpoint_parts_strip_chat_completions_suffix() {
let (host, prefix, query) = openai_compatible_endpoint_parts(
"https://router.huggingface.co/v1/chat/completions",
None,
)
.unwrap();
assert_eq!(host, "https://router.huggingface.co/v1");
assert_eq!(prefix, "");
assert!(query.is_empty());
}
#[test]
fn endpoint_parts_respect_explicit_base_path() {
let (host, prefix, query) = openai_compatible_endpoint_parts(
"https://router.huggingface.co",
Some("v1/chat/completions"),
)
.unwrap();
assert_eq!(host, "https://router.huggingface.co");
assert_eq!(prefix, "v1/");
assert!(query.is_empty());
}
#[tokio::test]
async fn custom_provider_returns_static_models_when_dynamic_models_disabled() {
let mut config = test_config();
config.requires_auth = false;
config.dynamic_models = Some(false);
config.models = vec![
ModelInfo::new("static-a".to_string(), 128000),
ModelInfo::new("static-b".to_string(), 128000),
];
let provider =
HuggingFaceProvider::from_custom_config(ModelConfig::new("static-a").unwrap(), config)
.unwrap();
assert_eq!(
provider.fetch_supported_models().await.unwrap(),
vec!["static-a".to_string(), "static-b".to_string()]
);
}
#[test]
fn custom_provider_requires_static_models_when_dynamic_models_disabled() {
let mut config = test_config();
config.requires_auth = false;
config.dynamic_models = Some(false);
let error = match HuggingFaceProvider::from_custom_config(
ModelConfig::new("model").unwrap(),
config,
) {
Ok(_) => panic!("expected dynamic_models: false without static models to fail"),
Err(error) => error,
};
assert_eq!(
error.to_string(),
"Provider 'custom_provider' has dynamic_models: false but no static models listed; at least one entry in `models` is required."
);
}
#[test]
fn custom_auth_method_respects_no_auth_config() {
let auth_method =
custom_auth_method_with_provider_token(false, Some("provider-token".to_string()))
.unwrap();
assert!(matches!(auth_method, AuthMethod::NoAuth));
}
#[test]
fn custom_auth_method_uses_provider_token_when_auth_is_required() {
let auth_method =
custom_auth_method_with_provider_token(true, Some("provider-token".to_string()))
.unwrap();
match auth_method {
AuthMethod::BearerToken(token) => assert_eq!(token, "provider-token"),
other => panic!("expected bearer token auth, got {other:?}"),
}
}
#[test]
fn custom_auth_method_uses_refresh_capable_auth_for_global_token() {
let auth_method = custom_auth_method_from_sources(true, None, || Ok(true)).unwrap();
assert!(matches!(auth_method, AuthMethod::Custom(_)));
}
#[test]
fn refreshable_huggingface_auth_method_uses_refresh_capable_auth() {
let auth_method = refreshable_huggingface_auth_method(|| Ok(true)).unwrap();
assert!(matches!(auth_method, AuthMethod::Custom(_)));
}
#[test]
fn refreshable_huggingface_auth_method_requires_configured_token() {
let error = refreshable_huggingface_auth_method(|| Ok(false)).unwrap_err();
assert_eq!(
error.to_string(),
"Hugging Face token is not configured. Sign in from Settings > Auth or configure HF_TOKEN."
);
}
#[test]
fn custom_auth_method_requires_global_token_when_auth_is_required() {
let error = custom_auth_method_from_sources(true, None, || Ok(false)).unwrap_err();
assert_eq!(
error.to_string(),
"Hugging Face token is not configured. Sign in from Settings > Auth or configure HF_TOKEN."
);
}
fn test_config() -> DeclarativeProviderConfig {
DeclarativeProviderConfig {
name: "custom_provider".to_string(),
engine: crate::config::declarative_providers::ProviderEngine::OpenAI,
display_name: "Custom Provider".to_string(),
description: None,
api_key_env: "CUSTOM_API_KEY".to_string(),
base_url: HUGGINGFACE_API_HOST.to_string(),
models: Vec::new(),
headers: None,
timeout_seconds: None,
supports_streaming: Some(true),
requires_auth: true,
catalog_provider_id: None,
base_path: None,
env_vars: None,
dynamic_models: None,
skip_canonical_filtering: false,
model_doc_link: None,
setup_steps: vec![],
fast_model: None,
preserves_thinking: true,
}
}
}
crates/goose/src/providers/huggingface_auth.rs
+896
View File
@@ -0,0 +1,896 @@
use crate::config::paths::Paths;
use crate::config::{Config, ConfigError};
use anyhow::{anyhow, Result};
use axum::{extract::Query, response::Html, routing::get, Router};
use base64::Engine;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::Digest;
use std::io;
use std::net::SocketAddr;
use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock};
use tokio::sync::{oneshot, Mutex as TokioMutex};
pub const HUGGINGFACE_PROVIDER_NAME: &str = "huggingface";
pub const HUGGINGFACE_DISPLAY_NAME: &str = "Hugging Face";
pub const HUGGINGFACE_TOKEN_SECRET_KEY: &str = "HF_TOKEN";
pub const HUGGINGFACE_OAUTH_TOKEN_NAME: &str = "OAuth token";
pub const HUGGINGFACE_OAUTH_CACHE_PATH: &str = "huggingface/oauth/tokens.json";
const AUTHORIZE_URL: &str = "https://huggingface.co/oauth/authorize";
const TOKEN_URL: &str = "https://huggingface.co/oauth/token";
const OAUTH_SCOPES: &str = "read-repos gated-repos inference-api";
const HUGGINGFACE_OAUTH_CLIENT_METADATA_URL: &str =
"https://goose-docs.ai/oauth/huggingface-client-metadata.json";
// This URI must match the redirect URI in the Hugging Face CIMD metadata.
const OAUTH_HOST: [u8; 4] = [127, 0, 0, 1];
const OAUTH_PORT: u16 = 17863;
const OAUTH_REDIRECT_PATH: &str = "/oauth/huggingface/callback";
const OAUTH_TIMEOUT_SECS: u64 = 300;
const HTML_AUTO_CLOSE_TIMEOUT_MS: u64 = 2000;
static HUGGINGFACE_OAUTH_MUTEX: LazyLock<TokioMutex<()>> = LazyLock::new(|| TokioMutex::new(()));
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HuggingFaceTokenData {
pub access_token: String,
#[serde(default)]
pub refresh_token: Option<String>,
#[serde(default)]
pub expires_at: Option<DateTime<Utc>>,
}
impl HuggingFaceTokenData {
pub fn is_expired(&self) -> bool {
self.expires_at
.is_some_and(|expires_at| expires_at <= Utc::now())
}
}
pub fn oauth_client_id() -> &'static str {
option_env!("GOOSE_HUGGINGFACE_OAUTH_CLIENT_ID")
.filter(|client_id| !client_id.trim().is_empty())
.unwrap_or(HUGGINGFACE_OAUTH_CLIENT_METADATA_URL)
}
pub fn oauth_cache_path() -> PathBuf {
Paths::in_config_dir(HUGGINGFACE_OAUTH_CACHE_PATH)
}
pub fn load_oauth_token() -> Option<HuggingFaceTokenData> {
load_oauth_token_from_path(&oauth_cache_path())
}
fn load_oauth_token_from_path(path: &Path) -> Option<HuggingFaceTokenData> {
let contents = std::fs::read_to_string(path).ok()?;
serde_json::from_str(&contents).ok()
}
pub fn has_oauth_token() -> bool {
load_oauth_token().is_some()
}
pub fn usable_oauth_token() -> Option<String> {
usable_oauth_token_from_path(&oauth_cache_path())
}
fn usable_oauth_token_from_path(path: &std::path::Path) -> Option<String> {
let token = load_oauth_token_from_path(path)?;
(!token.is_expired()).then_some(token.access_token)
}
pub fn has_usable_or_refreshable_oauth_token() -> bool {
has_usable_or_refreshable_oauth_token_from_path(&oauth_cache_path())
}
fn has_usable_or_refreshable_oauth_token_from_path(path: &std::path::Path) -> bool {
load_oauth_token_from_path(path).is_some_and(|token| {
!token.is_expired()
|| token
.refresh_token
.as_deref()
.is_some_and(|token| !token.is_empty())
})
}
pub fn has_configured_token() -> Result<bool> {
has_configured_token_from_sources(has_usable_or_refreshable_oauth_token(), hf_token_secret)
}
fn has_configured_token_from_sources(
has_oauth_token: bool,
secret_fallback: impl FnOnce() -> Result<Option<String>>,
) -> Result<bool> {
if has_oauth_token {
return Ok(true);
}
Ok(secret_fallback()?.is_some())
}
pub fn hf_token_secret() -> Result<Option<String>> {
match Config::global().get_secret::<String>(HUGGINGFACE_TOKEN_SECRET_KEY) {
Ok(token) => Ok(Some(token)),
Err(ConfigError::NotFound(_)) => Ok(None),
Err(error) => Err(error.into()),
}
}
pub fn resolve_token() -> Result<Option<String>> {
resolve_token_from_sources(None, usable_oauth_token(), hf_token_secret)
}
pub fn resolve_token_with_provider_token(provider_token: Option<String>) -> Result<Option<String>> {
resolve_token_from_sources(provider_token, usable_oauth_token(), hf_token_secret)
}
pub async fn resolve_token_async() -> Result<Option<String>> {
resolve_token_async_with_provider_token(None).await
}
pub async fn resolve_token_async_with_provider_token(
provider_token: Option<String>,
) -> Result<Option<String>> {
resolve_token_async_from_sources(
provider_token,
refreshed_or_usable_oauth_token_from_path(
&oauth_cache_path(),
oauth_client_id(),
TOKEN_URL,
),
hf_token_secret,
)
.await
}
async fn resolve_token_async_from_sources(
provider_token: Option<String>,
oauth_token: impl std::future::Future<Output = Result<Option<String>>>,
secret_fallback: impl FnOnce() -> Result<Option<String>>,
) -> Result<Option<String>> {
if provider_token.is_some() {
return Ok(provider_token);
}
match oauth_token.await {
Ok(Some(token)) => return Ok(Some(token)),
Ok(None) => {}
Err(refresh_error) => {
return match secret_fallback()? {
Some(token) => Ok(Some(token)),
None => Err(refresh_error),
};
}
}
secret_fallback()
}
fn resolve_token_from_sources(
provider_token: Option<String>,
oauth_token: Option<String>,
secret_fallback: impl FnOnce() -> Result<Option<String>>,
) -> Result<Option<String>> {
if provider_token.is_some() {
return Ok(provider_token);
}
if oauth_token.is_some() {
return Ok(oauth_token);
}
secret_fallback()
}
pub fn clear_oauth_token() -> Result<()> {
let path = oauth_cache_path();
if path.exists() {
std::fs::remove_file(path)?;
}
Ok(())
}
pub async fn configure_oauth() -> Result<()> {
let token_data = perform_loopback_oauth_flow(oauth_client_id()).await?;
save_oauth_token(token_data)
}
fn save_oauth_token(token_data: HuggingFaceTokenData) -> Result<()> {
let path = oauth_cache_path();
save_oauth_token_to_path(&path, &token_data)
}
fn save_oauth_token_to_path(path: &Path, token_data: &HuggingFaceTokenData) -> Result<()> {
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let contents = serde_json::to_string(&token_data)?;
std::fs::write(path, contents)?;
restrict_token_file_permissions(path)?;
Ok(())
}
#[cfg(unix)]
fn restrict_token_file_permissions(path: &Path) -> Result<()> {
use std::os::unix::fs::PermissionsExt;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o600))?;
Ok(())
}
#[cfg(not(unix))]
fn restrict_token_file_permissions(_path: &Path) -> Result<()> {
Ok(())
}
struct PkceChallenge {
verifier: String,
challenge: String,
}
fn generate_pkce() -> PkceChallenge {
let verifier = nanoid::nanoid!(64);
let digest = sha2::Sha256::digest(verifier.as_bytes());
let challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
PkceChallenge {
verifier,
challenge,
}
}
fn generate_state() -> String {
nanoid::nanoid!(32)
}
fn redirect_uri() -> String {
format!(
"http://{}.{}.{}.{}:{}{}",
OAUTH_HOST[0], OAUTH_HOST[1], OAUTH_HOST[2], OAUTH_HOST[3], OAUTH_PORT, OAUTH_REDIRECT_PATH
)
}
fn build_authorize_url(client_id: &str, pkce: &PkceChallenge, state: &str) -> Result<String> {
let redirect = redirect_uri();
let params = [
("response_type", "code"),
("client_id", client_id),
("redirect_uri", redirect.as_str()),
("scope", OAUTH_SCOPES),
("code_challenge", pkce.challenge.as_str()),
("code_challenge_method", "S256"),
("state", state),
];
let query = serde_urlencoded::to_string(params)?;
Ok(format!("{}?{}", AUTHORIZE_URL, query))
}
#[derive(Debug, Deserialize)]
struct TokenResponse {
access_token: String,
#[serde(default)]
refresh_token: Option<String>,
#[serde(default)]
expires_in: Option<i64>,
}
fn token_data_from_response(response: TokenResponse) -> HuggingFaceTokenData {
token_data_from_response_with_refresh_fallback(response, None)
}
fn token_data_from_response_with_refresh_fallback(
response: TokenResponse,
refresh_token_fallback: Option<String>,
) -> HuggingFaceTokenData {
HuggingFaceTokenData {
access_token: response.access_token,
refresh_token: response.refresh_token.or(refresh_token_fallback),
expires_at: response
.expires_in
.map(|secs| Utc::now() + chrono::Duration::seconds(secs)),
}
}
async fn exchange_code_for_tokens(
client_id: &str,
code: &str,
pkce: &PkceChallenge,
) -> Result<TokenResponse> {
let client = reqwest::Client::new();
let redirect = redirect_uri();
let params = [
("grant_type", "authorization_code"),
("code", code),
("redirect_uri", redirect.as_str()),
("client_id", client_id),
("code_verifier", pkce.verifier.as_str()),
];
let resp = client
.post(TOKEN_URL)
.header("Content-Type", "application/x-www-form-urlencoded")
.header("Accept", "application/json")
.form(&params)
.send()
.await?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(anyhow!(
"Hugging Face token exchange failed ({}): {}",
status,
text
));
}
Ok(resp.json().await?)
}
async fn refresh_access_token(
client_id: &str,
refresh_token: &str,
token_url: &str,
) -> Result<TokenResponse> {
let client = reqwest::Client::new();
let params = [
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", client_id),
];
let resp = client
.post(token_url)
.header("Content-Type", "application/x-www-form-urlencoded")
.header("Accept", "application/json")
.form(&params)
.send()
.await?;
if !resp.status().is_success() {
let status = resp.status();
let text = resp.text().await.unwrap_or_default();
return Err(anyhow!(
"Hugging Face token refresh failed ({}): {}",
status,
text
));
}
Ok(resp.json().await?)
}
async fn refreshed_or_usable_oauth_token_from_path(
path: &Path,
client_id: &str,
token_url: &str,
) -> Result<Option<String>> {
let Some(token) = load_oauth_token_from_path(path) else {
return Ok(None);
};
if !token.is_expired() {
return Ok(Some(token.access_token));
}
let Some(refresh_token) = token.refresh_token else {
return Ok(None);
};
let refreshed = refresh_access_token(client_id, &refresh_token, token_url).await?;
let refreshed =
token_data_from_response_with_refresh_fallback(refreshed, Some(refresh_token.clone()));
let access_token = refreshed.access_token.clone();
save_oauth_token_to_path(path, &refreshed)?;
Ok(Some(access_token))
}
const HTML_SUCCESS_TEMPLATE: &str = r#"<!doctype html>
<html>
<head>
<title>goose - Hugging Face Authorization Successful</title>
<script>setTimeout(() => window.close(), {timeout_ms});</script>
<style>
body {{
font-family: system-ui, -apple-system, sans-serif;
display: flex; justify-content: center; align-items: center;
height: 100vh; margin: 0; background: #171717; color: #fafafa;
}}
.container {{ text-align: center; padding: 2rem; }}
h1 {{ color: #ff9d00; margin-bottom: 1rem; }}
p {{ color: #c7c7c7; }}
</style>
</head>
<body>
<div class="container">
<h1>Authorization Successful</h1>
<p>You can close this window and return to goose.</p>
</div>
</body>
</html>"#;
fn html_success() -> String {
HTML_SUCCESS_TEMPLATE.replace("{timeout_ms}", &HTML_AUTO_CLOSE_TIMEOUT_MS.to_string())
}
fn html_error(error: &str) -> String {
let safe_error = v_htmlescape::escape_fmt(error);
format!(
r#"<!doctype html>
<html>
<head>
<title>goose - Hugging Face Authorization Failed</title>
<style>
body {{
font-family: system-ui, -apple-system, sans-serif;
display: flex; justify-content: center; align-items: center;
height: 100vh; margin: 0; background: #171717; color: #fafafa;
}}
.container {{ text-align: center; padding: 2rem; }}
h1 {{ color: #ff6b35; margin-bottom: 1rem; }}
p {{ color: #c7c7c7; }}
.error {{
color: #ffb199; font-family: monospace; margin-top: 1rem;
padding: 1rem; background: #3b180d; border-radius: 0.5rem;
}}
</style>
</head>
<body>
<div class="container">
<h1>Authorization Failed</h1>
<p>An error occurred during authorization.</p>
<div class="error">{}</div>
</div>
</body>
</html>"#,
safe_error
)
}
#[derive(Deserialize)]
struct CallbackParams {
code: Option<String>,
state: Option<String>,
error: Option<String>,
error_description: Option<String>,
}
fn oauth_callback_router(
expected_state: String,
tx: Arc<TokioMutex<Option<oneshot::Sender<Result<String>>>>>,
) -> Router {
Router::new().route(
OAUTH_REDIRECT_PATH,
get(move |Query(params): Query<CallbackParams>| {
let tx = tx.clone();
let expected = expected_state.clone();
async move {
if let Some(error) = params.error {
let msg = params.error_description.unwrap_or(error);
if let Some(sender) = tx.lock().await.take() {
let _ = sender.send(Err(anyhow!("{}", msg)));
}
return Html(html_error(&msg));
}
let code = match params.code {
Some(c) => c,
None => {
let msg = "Missing authorization code";
if let Some(sender) = tx.lock().await.take() {
let _ = sender.send(Err(anyhow!("{}", msg)));
}
return Html(html_error(msg));
}
};
if params.state.as_deref() != Some(&expected) {
let msg = "Invalid state - potential CSRF attack";
if let Some(sender) = tx.lock().await.take() {
let _ = sender.send(Err(anyhow!("{}", msg)));
}
return Html(html_error(msg));
}
if let Some(sender) = tx.lock().await.take() {
let _ = sender.send(Ok(code));
}
Html(html_success())
}
}),
)
}
async fn spawn_oauth_server(app: Router) -> Result<tokio::task::JoinHandle<()>> {
let addr = SocketAddr::from((OAUTH_HOST, OAUTH_PORT));
let listener = tokio::net::TcpListener::bind(addr).await.map_err(|e| {
if e.kind() == io::ErrorKind::AddrInUse {
anyhow!(
"Hugging Face OAuth callback server failed to bind to {}: port {} is already in use",
addr,
OAUTH_PORT
)
} else {
anyhow!(
"Hugging Face OAuth callback server failed to bind to {}: {}",
addr,
e
)
}
})?;
Ok(tokio::spawn(async move {
let server = axum::serve(listener, app);
let _ = server.await;
}))
}
struct ServerHandleGuard(Option<tokio::task::JoinHandle<()>>);
impl ServerHandleGuard {
fn new(handle: tokio::task::JoinHandle<()>) -> Self {
Self(Some(handle))
}
fn abort(&mut self) {
if let Some(handle) = self.0.take() {
handle.abort();
}
}
}
impl Drop for ServerHandleGuard {
fn drop(&mut self) {
self.abort();
}
}
async fn wait_for_oauth_code(rx: oneshot::Receiver<Result<String>>) -> Result<String> {
let code_result =
tokio::time::timeout(std::time::Duration::from_secs(OAUTH_TIMEOUT_SECS), rx).await;
code_result
.map_err(|_| anyhow!("Hugging Face OAuth flow timed out"))??
.map_err(|e| anyhow!("Hugging Face OAuth callback error: {}", e))
}
async fn perform_loopback_oauth_flow(client_id: &str) -> Result<HuggingFaceTokenData> {
let _guard = HUGGINGFACE_OAUTH_MUTEX.try_lock().map_err(|_| {
anyhow!("Another Hugging Face OAuth flow is already in progress; please try again later")
})?;
let pkce = generate_pkce();
let csrf_state = generate_state();
let auth_url = build_authorize_url(client_id, &pkce, &csrf_state)?;
let (tx, rx) = oneshot::channel::<Result<String>>();
let tx = Arc::new(TokioMutex::new(Some(tx)));
let app = oauth_callback_router(csrf_state.clone(), tx);
let server_handle = spawn_oauth_server(app).await?;
let mut server_guard = ServerHandleGuard::new(server_handle);
if webbrowser::open(&auth_url).is_err() {
tracing::info!(
"Please open this URL in your browser to authorize goose with Hugging Face:\n{}",
auth_url
);
}
let code_result = wait_for_oauth_code(rx).await;
server_guard.abort();
let code = code_result?;
let tokens = exchange_code_for_tokens(client_id, &code, &pkce).await?;
Ok(token_data_from_response(tokens))
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
use wiremock::matchers::{body_string_contains, method, path as request_path};
use wiremock::{Mock, MockServer, ResponseTemplate};
fn token_path(dir: &TempDir) -> PathBuf {
dir.path().join(HUGGINGFACE_OAUTH_CACHE_PATH)
}
fn with_token_path<T>(f: impl FnOnce(PathBuf) -> T) -> T {
let dir = TempDir::new().unwrap();
f(token_path(&dir))
}
#[test]
fn pkce_challenge_is_url_safe_base64_of_sha256_of_verifier() {
let pkce = generate_pkce();
assert_eq!(pkce.verifier.len(), 64);
assert_eq!(pkce.challenge.len(), 43);
assert!(!pkce.challenge.contains('='));
assert!(!pkce.challenge.contains('+'));
assert!(!pkce.challenge.contains('/'));
}
#[test]
fn authorize_url_contains_required_oauth_params() {
let pkce = PkceChallenge {
verifier: "v".repeat(64),
challenge: "challenge-fixture".to_string(),
};
let url = build_authorize_url("client-fixture", &pkce, "state-fixture").unwrap();
assert!(url.starts_with(AUTHORIZE_URL));
assert!(url.contains("client_id=client-fixture"));
assert!(url.contains("code_challenge=challenge-fixture"));
assert!(url.contains("code_challenge_method=S256"));
assert!(url.contains("state=state-fixture"));
assert!(url.contains("scope=read-repos"));
assert!(url.contains("gated-repos"));
assert!(url.contains("inference-api"));
}
#[test]
fn oauth_client_id_defaults_to_cimd_metadata_url() {
if option_env!("GOOSE_HUGGINGFACE_OAUTH_CLIENT_ID").is_none() {
assert_eq!(oauth_client_id(), HUGGINGFACE_OAUTH_CLIENT_METADATA_URL);
}
}
#[test]
fn redirect_uri_matches_huggingface_cimd_metadata() {
assert_eq!(
redirect_uri(),
"http://127.0.0.1:17863/oauth/huggingface/callback"
);
}
#[test]
fn token_data_from_response_stores_expires_in_as_expires_at() {
let token_data = token_data_from_response(TokenResponse {
access_token: "token".to_string(),
refresh_token: None,
expires_in: Some(60),
});
let expires_at = token_data.expires_at.unwrap();
assert!(expires_at > Utc::now());
assert!(expires_at <= Utc::now() + chrono::Duration::seconds(60));
}
#[tokio::test]
async fn expired_oauth_token_refreshes_with_cached_refresh_token() {
let dir = TempDir::new().unwrap();
let path = token_path(&dir);
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "expired".to_string(),
refresh_token: Some("refresh".to_string()),
expires_at: Some(Utc::now() - chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(request_path("/"))
.and(body_string_contains("grant_type=refresh_token"))
.and(body_string_contains("refresh_token=refresh"))
.and(body_string_contains("client_id=client-fixture"))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"access_token": "refreshed",
"expires_in": 60
})))
.mount(&server)
.await;
let token =
refreshed_or_usable_oauth_token_from_path(&path, "client-fixture", &server.uri())
.await
.unwrap();
assert_eq!(token.as_deref(), Some("refreshed"));
let saved = load_oauth_token_from_path(&path).unwrap();
assert_eq!(saved.access_token, "refreshed");
assert_eq!(saved.refresh_token.as_deref(), Some("refresh"));
assert!(saved.expires_at.unwrap() > Utc::now());
}
#[test]
fn usable_oauth_token_skips_expired_token() {
with_token_path(|path| {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "expired".to_string(),
refresh_token: None,
expires_at: Some(Utc::now() - chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
assert_eq!(usable_oauth_token_from_path(&path), None);
});
}
#[test]
fn usable_oauth_token_returns_unexpired_token() {
with_token_path(|path| {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "valid".to_string(),
refresh_token: None,
expires_at: Some(Utc::now() + chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
assert_eq!(
usable_oauth_token_from_path(&path).as_deref(),
Some("valid")
);
});
}
#[test]
fn has_usable_or_refreshable_oauth_token_accepts_unexpired_token() {
with_token_path(|path| {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "valid".to_string(),
refresh_token: None,
expires_at: Some(Utc::now() + chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
assert!(has_usable_or_refreshable_oauth_token_from_path(&path));
});
}
#[test]
fn has_usable_or_refreshable_oauth_token_accepts_expired_refreshable_token() {
with_token_path(|path| {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "expired".to_string(),
refresh_token: Some("refresh".to_string()),
expires_at: Some(Utc::now() - chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
assert!(has_usable_or_refreshable_oauth_token_from_path(&path));
});
}
#[test]
fn has_usable_or_refreshable_oauth_token_rejects_expired_unrefreshable_token() {
with_token_path(|path| {
std::fs::create_dir_all(path.parent().unwrap()).unwrap();
std::fs::write(
&path,
serde_json::to_string(&HuggingFaceTokenData {
access_token: "expired".to_string(),
refresh_token: None,
expires_at: Some(Utc::now() - chrono::Duration::minutes(1)),
})
.unwrap(),
)
.unwrap();
assert!(!has_usable_or_refreshable_oauth_token_from_path(&path));
});
}
#[test]
fn has_configured_token_accepts_oauth_without_secret_lookup() {
let configured = has_configured_token_from_sources(true, || {
panic!("secret store should not be queried when OAuth is configured")
})
.unwrap();
assert!(configured);
}
#[test]
fn has_configured_token_accepts_secret_fallback() {
let configured =
has_configured_token_from_sources(false, || Ok(Some("hf-token".to_string()))).unwrap();
assert!(configured);
}
#[test]
fn has_configured_token_rejects_missing_oauth_and_secret() {
let configured = has_configured_token_from_sources(false, || Ok(None)).unwrap();
assert!(!configured);
}
#[cfg(unix)]
#[test]
fn save_oauth_token_restricts_file_permissions() {
use std::os::unix::fs::PermissionsExt;
with_token_path(|path| {
save_oauth_token_to_path(
&path,
&HuggingFaceTokenData {
access_token: "saved".to_string(),
refresh_token: None,
expires_at: None,
},
)
.unwrap();
let mode = std::fs::metadata(&path).unwrap().permissions().mode() & 0o777;
assert_eq!(mode, 0o600);
});
}
#[test]
fn resolver_prefers_provider_token_over_oauth() {
let token = resolve_token_from_sources(
Some("api-key".to_string()),
Some("oauth".to_string()),
|| panic!("secret store should not be queried when provider token is usable"),
)
.unwrap();
assert_eq!(token.as_deref(), Some("api-key"));
}
#[test]
fn resolver_uses_oauth_before_secret_store() {
let token = resolve_token_from_sources(None, Some("oauth".to_string()), || {
panic!("secret store should not be queried when OAuth is usable")
})
.unwrap();
assert_eq!(token.as_deref(), Some("oauth"));
}
#[test]
fn resolver_uses_secret_store_when_no_provider_token_or_oauth_exists() {
let token = resolve_token_from_sources(None, None, || Ok(Some("secret-store".to_string())))
.unwrap();
assert_eq!(token.as_deref(), Some("secret-store"));
}
#[tokio::test]
async fn async_resolver_uses_secret_fallback_when_oauth_refresh_fails() {
let token = resolve_token_async_from_sources(
None,
async { Err(anyhow::anyhow!("refresh token revoked")) },
|| Ok(Some("secret-store".to_string())),
)
.await
.unwrap();
assert_eq!(token.as_deref(), Some("secret-store"));
}
#[tokio::test]
async fn async_resolver_reports_refresh_error_without_secret_fallback() {
let error = resolve_token_async_from_sources(
None,
async { Err(anyhow::anyhow!("refresh token revoked")) },
|| Ok(None),
)
.await
.unwrap_err();
assert_eq!(error.to_string(), "refresh token revoked");
}
}
crates/goose/src/providers/init.rs
+26
View File
@@ -27,6 +27,7 @@ use super::{
gemini_oauth::GeminiOAuthProvider,
githubcopilot::GithubCopilotProvider,
google::GoogleProvider,
huggingface::HuggingFaceProvider,
kimicode::KimiCodeProvider,
litellm::LiteLLMProvider,
nanogpt::NanoGptProvider,
@@ -76,6 +77,7 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
registry.register::<GeminiOAuthProvider>(true);
registry.register::<GithubCopilotProvider>(false);
registry.register::<GoogleProvider>(true);
registry.register::<HuggingFaceProvider>(true);
registry.register::<KimiCodeProvider>(true);
registry.register::<LiteLLMProvider>(false);
registry.register::<NanoGptProvider>(true);
@@ -111,10 +113,18 @@ async fn init_registry() -> RwLock<ProviderRegistry> {
"chatgpt_codex",
Arc::new(|| Box::pin(ChatGptCodexProvider::cleanup())),
);
registry.set_cleanup(
"gemini_oauth",
Arc::new(|| Box::pin(GeminiOAuthProvider::cleanup())),
);
registry.set_cleanup(
"xai_oauth",
Arc::new(|| Box::pin(XaiOAuthProvider::cleanup())),
);
registry.set_cleanup(
"huggingface",
Arc::new(|| Box::pin(HuggingFaceProvider::cleanup())),
);
if let Err(e) = load_custom_providers_into_registry(&mut registry) {
tracing::warn!("Failed to load custom providers: {}", e);
@@ -261,6 +271,22 @@ mod tests {
assert!(!endpoint.secret, "Endpoint should not be secret");
}
#[tokio::test]
async fn test_huggingface_provider_registry_wiring() {
let huggingface = get_from_registry("huggingface")
.await
.expect("huggingface provider should be registered");
let meta = huggingface.metadata();
assert_eq!(huggingface.provider_type(), ProviderType::Preferred);
assert_eq!(meta.display_name, "Hugging Face");
assert_eq!(meta.default_model, "Qwen/Qwen3-Coder-480B-A35B-Instruct");
assert!(meta
.config_keys
.iter()
.any(|key| key.name == "HF_TOKEN" && key.secret));
}
#[tokio::test]
async fn test_nvidia_declarative_provider_registry_wiring() {
let nvidia = get_from_registry("nvidia")
crates/goose/src/providers/local_inference/hf_models.rs
+55 -8
View File
@@ -1,6 +1,8 @@
use anyhow::{bail, Result};
use serde::{Deserialize, Serialize};
use crate::providers::huggingface_auth;
use utoipa::ToSchema;
const HF_API_BASE: &str = "https://huggingface.co/api/models";
@@ -245,6 +247,26 @@ fn build_download_url(repo_id: &str, filename: &str) -> String {
format!("{}/{}/resolve/main/{}", HF_DOWNLOAD_BASE, repo_id, filename)
}
pub fn hf_authorization_header(token: Option<&str>) -> Option<String> {
token
.filter(|token| !token.is_empty())
.map(|token| format!("Bearer {}", token))
}
fn apply_hf_auth(request: reqwest::RequestBuilder, token: Option<&str>) -> reqwest::RequestBuilder {
if let Some(header) = hf_authorization_header(token) {
request.header("Authorization", header)
} else {
request
}
}
async fn optional_hf_token(
token: impl std::future::Future<Output = Result<Option<String>>>,
) -> Option<String> {
token.await.ok().flatten()
}
fn parent_components(filename: &str) -> Vec<&str> {
filename.rsplit_once('/').map_or(Vec::new(), |(parent, _)| {
parent.split('/').filter(|part| !part.is_empty()).collect()
@@ -404,13 +426,13 @@ fn group_into_variants(repo_id: &str, files: Vec<HfApiSibling>) -> Vec<HfQuantVa
pub async fn search_gguf_models(query: &str, limit: usize) -> Result<Vec<HfModelInfo>> {
let client = reqwest::Client::new();
let token = optional_hf_token(huggingface_auth::resolve_token_async()).await;
let url = format!(
"{}?search={}&filter=gguf&sort=downloads&direction=-1&limit={}",
HF_API_BASE, query, limit
);
let response = client
.get(&url)
let response = apply_hf_auth(client.get(&url), token.as_deref())
.header("User-Agent", "goose-ai-agent")
.send()
.await?;
@@ -469,10 +491,10 @@ pub async fn search_gguf_models(query: &str, limit: usize) -> Result<Vec<HfModel
/// Fetch GGUF files for a repo and return them grouped by quantization.
pub async fn get_repo_gguf_variants(repo_id: &str) -> Result<Vec<HfQuantVariant>> {
let client = reqwest::Client::new();
let token = optional_hf_token(huggingface_auth::resolve_token_async()).await;
let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id);
let response = client
.get(&url)
let response = apply_hf_auth(client.get(&url), token.as_deref())
.header("User-Agent", "goose-ai-agent")
.send()
.await?;
@@ -494,10 +516,10 @@ pub async fn get_repo_gguf_variants(repo_id: &str) -> Result<Vec<HfQuantVariant>
/// Fetch raw GGUF files (kept for resolve_model_spec).
pub async fn get_repo_gguf_files(repo_id: &str) -> Result<Vec<HfGgufFile>> {
let client = reqwest::Client::new();
let token = optional_hf_token(huggingface_auth::resolve_token_async()).await;
let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id);
let response = client
.get(&url)
let response = apply_hf_auth(client.get(&url), token.as_deref())
.header("User-Agent", "goose-ai-agent")
.send()
.await?;
@@ -556,9 +578,9 @@ pub async fn resolve_model_spec_full(spec: &str) -> Result<(String, ResolvedMode
let (repo_id, quant) = parse_model_spec(spec)?;
let client = reqwest::Client::new();
let token = optional_hf_token(huggingface_auth::resolve_token_async()).await;
let url = format!("{}/{}?blobs=true", HF_API_BASE, repo_id);
let response = client
.get(&url)
let response = apply_hf_auth(client.get(&url), token.as_deref())
.header("User-Agent", "goose-ai-agent")
.send()
.await?;
@@ -739,6 +761,16 @@ mod tests {
assert_eq!(parse_quantization("random-name.gguf"), "unknown");
}
#[test]
fn test_hf_authorization_header() {
assert_eq!(
hf_authorization_header(Some("hf_test")).as_deref(),
Some("Bearer hf_test")
);
assert_eq!(hf_authorization_header(Some("")), None);
assert_eq!(hf_authorization_header(None), None);
}
#[test]
fn test_parse_quantization_with_directory() {
assert_eq!(
@@ -926,6 +958,21 @@ mod tests {
assert_eq!(mmproj.quantization, "BF16");
}
#[tokio::test]
async fn optional_hf_token_returns_resolved_token() {
let token = optional_hf_token(async { Ok(Some("token".to_string())) }).await;
assert_eq!(token.as_deref(), Some("token"));
}
#[tokio::test]
async fn optional_hf_token_ignores_resolution_errors() {
let token =
optional_hf_token(async { Err(anyhow::anyhow!("refresh token revoked")) }).await;
assert_eq!(token, None);
}
#[test]
fn test_select_best_mmproj_prefers_bf16_over_f16_tie() {
let files = vec![
crates/goose/src/providers/mod.rs
+2
View File
@@ -31,6 +31,8 @@ pub mod gemini_oauth;
pub mod githubcopilot;
pub mod google;
pub mod http_status;
pub mod huggingface;
pub mod huggingface_auth;
mod init;
pub mod inventory;
pub mod kimicode;
crates/goose/src/providers/openai_compatible.rs
+62 -4
View File
@@ -11,13 +11,15 @@ use tokio_util::codec::{FramedRead, LinesCodec};
use tokio_util::io::StreamReader;
use super::api_client::ApiClient;
use super::base::{MessageStream, Provider};
use super::base::{stream_from_single_message, MessageStream, Provider, ProviderUsage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{ImageFormat, RequestLog};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::formats::openai::{create_request, response_to_streaming_message};
use crate::providers::formats::openai::{
create_request, get_usage, response_to_message, response_to_streaming_message,
};
use crate::providers::formats::openai_responses::responses_api_to_streaming_message;
use rmcp::model::Tool;
@@ -28,6 +30,7 @@ pub struct OpenAiCompatibleProvider {
model: ModelConfig,
/// Path prefix prepended to `chat/completions` (e.g. `"deployments/{name}/"` for Azure).
completions_prefix: String,
supports_streaming: bool,
}
impl OpenAiCompatibleProvider {
@@ -42,9 +45,15 @@ impl OpenAiCompatibleProvider {
api_client,
model,
completions_prefix,
supports_streaming: true,
}
}
pub fn with_supports_streaming(mut self, supports_streaming: bool) -> Self {
self.supports_streaming = supports_streaming;
self
}
fn build_request(
&self,
model_config: &ModelConfig,
@@ -110,7 +119,13 @@ impl Provider for OpenAiCompatibleProvider {
messages: &[Message],
tools: &[Tool],
) -> Result<MessageStream, ProviderError> {
let payload = self.build_request(model_config, system, messages, tools, true)?;
let payload = self.build_request(
model_config,
system,
messages,
tools,
self.supports_streaming,
)?;
let mut log = RequestLog::start(model_config, &payload)?;
let completions_path = format!("{}chat/completions", self.completions_prefix);
@@ -127,7 +142,27 @@ impl Provider for OpenAiCompatibleProvider {
let _ = log.error(e);
})?;
stream_openai_compat(response, log)
if self.supports_streaming {
stream_openai_compat(response, log)
} else {
let json: serde_json::Value = response.json().await.map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse JSON: {}", e))
})?;
let message = response_to_message(&json).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to parse message: {}", e))
})?;
let usage_data = get_usage(json.get("usage").unwrap_or(&serde_json::Value::Null));
let usage = ProviderUsage::new(model_config.model_name.clone(), usage_data);
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
Some(&usage.usage),
)?;
Ok(stream_from_single_message(message, usage))
}
}
}
@@ -190,6 +225,7 @@ pub fn stream_responses_compat(
#[cfg(test)]
mod tests {
use super::*;
use crate::model::ModelConfig;
use serde_json::json;
use test_case::test_case;
@@ -262,4 +298,26 @@ mod tests {
"Expected {expected_variant}, got error: {err:?}"
);
}
#[test]
fn build_request_respects_non_streaming_mode() {
let provider = OpenAiCompatibleProvider::new(
"test".to_string(),
ApiClient::new(
"http://localhost".to_string(),
super::super::api_client::AuthMethod::NoAuth,
)
.unwrap(),
ModelConfig::new_or_fail("test-model"),
String::new(),
)
.with_supports_streaming(false);
let payload = provider
.build_request(&provider.model, "", &[], &[], provider.supports_streaming)
.unwrap();
assert_eq!(payload.get("stream"), None);
assert_eq!(payload.get("stream_options"), None);
}
}
crates/goose/src/providers/provider_registry.rs
+103 -6
View File
@@ -161,6 +161,53 @@ impl ProviderRegistry {
P: ProviderDef + 'static,
F: Fn(ModelConfig) -> Result<P::Provider> + Send + Sync + 'static,
G: Fn() -> Result<InventoryIdentityInput> + Send + Sync + 'static,
{
self.register_with_name_impl::<P, F, G>(
config,
provider_type,
supports_inventory_refresh,
constructor,
inventory_identity,
None,
);
}
pub fn register_with_name_and_inventory_configured<P, F, G, H>(
&mut self,
config: &DeclarativeProviderConfig,
provider_type: ProviderType,
supports_inventory_refresh: bool,
constructor: F,
inventory_identity: G,
inventory_configured: H,
) where
P: ProviderDef + 'static,
F: Fn(ModelConfig) -> Result<P::Provider> + Send + Sync + 'static,
G: Fn() -> Result<InventoryIdentityInput> + Send + Sync + 'static,
H: Fn() -> bool + Send + Sync + 'static,
{
self.register_with_name_impl::<P, F, G>(
config,
provider_type,
supports_inventory_refresh,
constructor,
inventory_identity,
Some(Arc::new(inventory_configured)),
);
}
fn register_with_name_impl<P, F, G>(
&mut self,
config: &DeclarativeProviderConfig,
provider_type: ProviderType,
supports_inventory_refresh: bool,
constructor: F,
inventory_identity: G,
inventory_configured: Option<ProviderInventoryConfiguredResolver>,
) where
P: ProviderDef + 'static,
F: Fn(ModelConfig) -> Result<P::Provider> + Send + Sync + 'static,
G: Fn() -> Result<InventoryIdentityInput> + Send + Sync + 'static,
{
let base_metadata = P::metadata();
let description = config
@@ -243,6 +290,12 @@ impl ProviderRegistry {
model_selection_hint: None,
};
let inventory_config_keys = custom_metadata.config_keys.clone();
let default_inventory_configured = Arc::new(move || {
super::inventory::default_inventory_configured(
&inventory_config_keys,
crate::config::Config::global(),
)
});
self.entries.insert(
config.name.clone(),
@@ -256,12 +309,7 @@ impl ProviderRegistry {
})
}),
inventory_identity: Arc::new(inventory_identity),
inventory_configured: Arc::new(move || {
super::inventory::default_inventory_configured(
&inventory_config_keys,
crate::config::Config::global(),
)
}),
inventory_configured: inventory_configured.unwrap_or(default_inventory_configured),
cleanup: None,
provider_type,
supports_inventory_refresh,
@@ -308,3 +356,52 @@ impl ProviderRegistry {
self.entries.retain(|name, _| !name.starts_with("custom_"));
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::config::declarative_providers::ProviderEngine;
use crate::providers::openai::OpenAiProvider;
fn test_config() -> DeclarativeProviderConfig {
DeclarativeProviderConfig {
name: "custom_hf".to_string(),
engine: ProviderEngine::OpenAI,
display_name: "Custom HF".to_string(),
description: None,
api_key_env: String::new(),
base_url: "https://router.huggingface.co/v1".to_string(),
models: vec![ModelInfo::new("test-model", 128_000)],
headers: None,
timeout_seconds: None,
supports_streaming: Some(true),
requires_auth: true,
catalog_provider_id: Some("huggingface".to_string()),
base_path: None,
env_vars: None,
dynamic_models: None,
skip_canonical_filtering: false,
model_doc_link: None,
setup_steps: vec![],
fast_model: None,
preserves_thinking: false,
}
}
#[test]
fn register_with_name_can_override_inventory_configured() {
let mut registry = ProviderRegistry::new();
registry.register_with_name_and_inventory_configured::<OpenAiProvider, _, _, _>(
&test_config(),
ProviderType::Declarative,
false,
|_| unreachable!("constructor is not used by this test"),
|| Ok(InventoryIdentityInput::new("custom_hf", "huggingface")),
|| false,
);
let entry = registry.entries.get("custom_hf").unwrap();
assert!(!entry.inventory_configured());
}
}