overhaul provider inventory and agent/model selection (#8652)

Signed-off-by: Bradley Axen <baxen@squareup.com>
This commit is contained in:
Bradley Axen
2026-04-20 15:00:17 -07:00
committed by GitHub
parent 3d582943fd
commit 8eda6fdabc
70 changed files with 5321 additions and 2123 deletions
@@ -416,6 +416,11 @@ mod tests {
use std::sync::Arc;
let tmp_dir = tempfile::tempdir().unwrap();
let temp_root = tmp_dir.path().display().to_string();
let _guard = env_lock::lock_env([
("HOME", Some(temp_root.as_str())),
("GOOSE_PATH_ROOT", Some(temp_root.as_str())),
]);
let session_manager = Arc::new(SessionManager::new(tmp_dir.path().to_path_buf()));
let session = session_manager
.create_session(
@@ -2,6 +2,7 @@ use crate::config::paths::Paths;
use crate::config::Config;
use crate::providers::anthropic::AnthropicProvider;
use crate::providers::base::{ModelInfo, ProviderType};
use crate::providers::inventory::declarative_inventory_identity;
use crate::providers::ollama::OllamaProvider;
use crate::providers::openai::OpenAiProvider;
use anyhow::Result;
@@ -460,38 +461,59 @@ pub fn register_declarative_provider(
match config.engine {
ProviderEngine::OpenAI => {
let captured = config.clone();
registry.register_with_name::<OpenAiProvider, _>(
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)
},
);
}
ProviderEngine::Ollama => {
let captured = config.clone();
registry.register_with_name::<OllamaProvider, _>(
let identity_config = config.clone();
registry.register_with_name::<OllamaProvider, _, _>(
&config,
provider_type,
config.dynamic_models.unwrap_or(false),
move |model| {
let mut cfg = captured.clone();
resolve_config(&mut cfg)?;
OllamaProvider::from_custom_config(model, cfg)
},
move || {
let mut cfg = identity_config.clone();
resolve_config(&mut cfg)?;
declarative_inventory_identity(&cfg)
},
);
}
ProviderEngine::Anthropic => {
let captured = config.clone();
registry.register_with_name::<AnthropicProvider, _>(
let identity_config = config.clone();
registry.register_with_name::<AnthropicProvider, _, _>(
&config,
provider_type,
config.dynamic_models.unwrap_or(false),
move |model| {
let mut cfg = captured.clone();
resolve_config(&mut cfg)?;
AnthropicProvider::from_custom_config(model, cfg)
},
move || {
let mut cfg = identity_config.clone();
resolve_config(&mut cfg)?;
declarative_inventory_identity(&cfg)
},
);
}
}
+18
View File
@@ -0,0 +1,18 @@
use crate::config::search_path::SearchPaths;
use crate::providers::inventory::InventoryIdentityInput;
use anyhow::Result;
use std::path::PathBuf;
pub fn acp_adapter_installed(command: &str) -> bool {
resolve_acp_command(command).is_ok()
}
pub fn acp_inventory_identity(provider_id: &str, command: &str) -> Result<InventoryIdentityInput> {
let resolved_command = resolve_acp_command(command)?;
Ok(InventoryIdentityInput::new(provider_id, provider_id)
.with_public("command", resolved_command.display().to_string()))
}
fn resolve_acp_command(command: &str) -> Result<PathBuf> {
SearchPaths::builder().with_npm().resolve(command)
}
+21 -4
View File
@@ -9,7 +9,9 @@ use crate::acp::{
use crate::config::search_path::SearchPaths;
use crate::config::{Config, GooseMode};
use crate::model::ModelConfig;
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
use crate::providers::base::{ProviderDef, ProviderMetadata};
use crate::providers::inventory::InventoryIdentityInput;
const AMP_ACP_PROVIDER_NAME: &str = "amp-acp";
const AMP_ACP_DOC_URL: &str = "https://ampcode.com";
@@ -37,6 +39,7 @@ impl ProviderDef for AmpAcpProvider {
"Set in your goose config file (`~/.config/goose/config.yaml` on macOS/Linux):\n GOOSE_PROVIDER: amp-acp\n GOOSE_MODEL: current",
"Restart goose for changes to take effect",
])
.with_model_selection_hint("Use the Amp CLI to configure models")
}
fn from_env(
@@ -49,10 +52,12 @@ impl ProviderDef for AmpAcpProvider {
let goose_mode = config.get_goose_mode().unwrap_or(GooseMode::Auto);
let mode_mapping = HashMap::from([
(GooseMode::Auto, "auto".to_string()),
(GooseMode::Approve, "approve".to_string()),
(GooseMode::SmartApprove, "smart-approve".to_string()),
(GooseMode::Chat, "chat".to_string()),
// "bypass" skips confirmations, closest to autonomous mode.
(GooseMode::Auto, "bypass".to_string()),
// "default" prompts before risky actions.
(GooseMode::Approve, "default".to_string()),
(GooseMode::SmartApprove, "default".to_string()),
(GooseMode::Chat, "default".to_string()),
]);
let provider_config = AcpProviderConfig {
@@ -71,4 +76,16 @@ impl ProviderDef for AmpAcpProvider {
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
})
}
fn supports_inventory_refresh() -> bool {
false
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
acp_inventory_identity(AMP_ACP_PROVIDER_NAME, AMP_ACP_BINARY)
}
fn inventory_configured() -> bool {
acp_adapter_installed(AMP_ACP_BINARY)
}
}
+28
View File
@@ -14,6 +14,7 @@ use super::errors::ProviderError;
use super::formats::anthropic::{
create_request, response_to_streaming_message, thinking_type, ThinkingType,
};
use super::inventory::{config_secret_value, serialize_string_map, InventoryIdentityInput};
use super::openai_compatible::handle_status_openai_compat;
use super::openai_compatible::map_http_error_to_provider_error;
use super::retry::ProviderRetry;
@@ -235,6 +236,33 @@ impl ProviderDef for AnthropicProvider {
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(Self::from_env(model))
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
let config = crate::config::Config::global();
let mut identity =
InventoryIdentityInput::new(ANTHROPIC_PROVIDER_NAME, ANTHROPIC_PROVIDER_NAME)
.with_public(
"host",
config
.get_param::<String>("ANTHROPIC_HOST")
.unwrap_or_else(|_| "https://api.anthropic.com".to_string()),
);
if let Some(api_key) = config_secret_value(config, "ANTHROPIC_API_KEY") {
identity = identity.with_secret("api_key", api_key);
}
if let Ok(headers) = config
.get_secret::<std::collections::HashMap<String, String>>("ANTHROPIC_CUSTOM_HEADERS")
{
identity = identity.with_secret("headers", serialize_string_map(&headers)?);
}
Ok(identity)
}
}
#[async_trait]
+45 -5
View File
@@ -6,9 +6,10 @@ use serde::{Deserialize, Serialize};
use super::canonical::{map_to_canonical_model, CanonicalModelRegistry};
use super::errors::ProviderError;
use super::inventory::{default_inventory_identity, InventoryIdentityInput};
use super::retry::RetryConfig;
use crate::config::base::ConfigValue;
use crate::config::{ExtensionConfig, GooseMode};
use crate::config::{Config, ExtensionConfig, GooseMode};
use crate::conversation::message::{Message, MessageContent};
use crate::conversation::Conversation;
use crate::model::ModelConfig;
@@ -179,6 +180,9 @@ pub struct ProviderMetadata {
/// step-by-step instructions for set up providers eg: api key
#[serde(default)]
pub setup_steps: Vec<String>,
/// Hint shown in the model picker when this provider manages its own model selection.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub model_selection_hint: Option<String>,
}
impl ProviderMetadata {
@@ -212,6 +216,7 @@ impl ProviderMetadata {
model_doc_link: model_doc_link.to_string(),
config_keys,
setup_steps: vec![],
model_selection_hint: None,
}
}
@@ -233,6 +238,7 @@ impl ProviderMetadata {
model_doc_link: model_doc_link.to_string(),
config_keys,
setup_steps: vec![],
model_selection_hint: None,
}
}
@@ -246,6 +252,7 @@ impl ProviderMetadata {
model_doc_link: "".to_string(),
config_keys: vec![],
setup_steps: vec![],
model_selection_hint: None,
}
}
@@ -253,6 +260,11 @@ impl ProviderMetadata {
self.setup_steps = steps.into_iter().map(|s| s.to_string()).collect();
self
}
pub fn with_model_selection_hint(mut self, hint: &str) -> Self {
self.model_selection_hint = Some(hint.to_string());
self
}
}
/// Configuration key metadata for provider setup
@@ -492,6 +504,34 @@ pub trait ProviderDef: Send + Sync {
) -> BoxFuture<'static, Result<Self::Provider>>
where
Self: Sized;
fn supports_inventory_refresh() -> bool
where
Self: Sized,
{
false
}
fn inventory_identity() -> Result<InventoryIdentityInput>
where
Self: Sized,
{
let metadata = Self::metadata();
Ok(default_inventory_identity(
&metadata.name,
&metadata.name,
&metadata.config_keys,
Config::global(),
))
}
fn inventory_configured() -> bool
where
Self: Sized,
{
let metadata = Self::metadata();
super::inventory::default_inventory_configured(&metadata.config_keys, Config::global())
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
@@ -588,7 +628,7 @@ pub trait Provider: Send + Sync {
false
}
/// Fetch models filtered by canonical registry and usability
/// Fetch inventory models filtered by canonical registry and usability.
async fn fetch_recommended_models(&self) -> Result<Vec<String>, ProviderError> {
let all_models = self.fetch_supported_models().await?;
@@ -637,15 +677,15 @@ pub trait Provider: Send + Sync {
(None, None) => a.0.cmp(&b.0),
});
let recommended_models: Vec<String> = models_with_dates
let inventory_models: Vec<String> = models_with_dates
.into_iter()
.map(|(name, _)| name)
.collect();
if recommended_models.is_empty() {
if inventory_models.is_empty() {
Ok(all_models)
} else {
Ok(recommended_models)
Ok(inventory_models)
}
}
+14
View File
@@ -9,7 +9,9 @@ use crate::acp::{
use crate::config::search_path::SearchPaths;
use crate::config::{Config, GooseMode};
use crate::model::ModelConfig;
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
use crate::providers::base::{ProviderDef, ProviderMetadata};
use crate::providers::inventory::InventoryIdentityInput;
const CLAUDE_ACP_PROVIDER_NAME: &str = "claude-acp";
const CLAUDE_ACP_DOC_URL: &str = "https://github.com/zed-industries/claude-agent-acp";
@@ -78,4 +80,16 @@ impl ProviderDef for ClaudeAcpProvider {
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
})
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
acp_inventory_identity(CLAUDE_ACP_PROVIDER_NAME, CLAUDE_ACP_BINARY)
}
fn inventory_configured() -> bool {
acp_adapter_installed(CLAUDE_ACP_BINARY)
}
}
+14
View File
@@ -9,7 +9,9 @@ use crate::acp::{
use crate::config::search_path::SearchPaths;
use crate::config::{Config, GooseMode};
use crate::model::ModelConfig;
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
use crate::providers::base::{ProviderDef, ProviderMetadata};
use crate::providers::inventory::InventoryIdentityInput;
const CODEX_ACP_PROVIDER_NAME: &str = "codex-acp";
const CODEX_ACP_DOC_URL: &str = "https://github.com/zed-industries/codex-acp";
@@ -98,6 +100,18 @@ impl ProviderDef for CodexAcpProvider {
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
})
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
acp_inventory_identity(CODEX_ACP_PROVIDER_NAME, CODEX_ACP_PROVIDER_NAME)
}
fn inventory_configured() -> bool {
acp_adapter_installed(CODEX_ACP_PROVIDER_NAME)
}
}
// Codex sandbox scope determines what needs approval: operations within the
+14
View File
@@ -9,7 +9,9 @@ use crate::acp::{
use crate::config::search_path::SearchPaths;
use crate::config::{Config, GooseMode};
use crate::model::ModelConfig;
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
use crate::providers::base::{ProviderDef, ProviderMetadata};
use crate::providers::inventory::InventoryIdentityInput;
const COPILOT_ACP_PROVIDER_NAME: &str = "copilot-acp";
const COPILOT_ACP_DOC_URL: &str = "https://github.com/github/copilot-cli";
@@ -84,4 +86,16 @@ impl ProviderDef for CopilotAcpProvider {
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
})
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
acp_inventory_identity(COPILOT_ACP_PROVIDER_NAME, COPILOT_ACP_BINARY)
}
fn inventory_configured() -> bool {
acp_adapter_installed(COPILOT_ACP_BINARY)
}
}
+4
View File
@@ -340,6 +340,10 @@ impl ProviderDef for DatabricksProvider {
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(Self::from_env(model))
}
fn supports_inventory_refresh() -> bool {
true
}
}
#[async_trait]
@@ -1100,12 +1100,16 @@ mod tests {
fn test_create_request_enabled_thinking_with_budget() -> anyhow::Result<()> {
let _guard = env_lock::lock_env([
("CLAUDE_THINKING_TYPE", None::<&str>),
("CLAUDE_THINKING_ENABLED", Some("1")),
("CLAUDE_THINKING_ENABLED", None::<&str>),
("CLAUDE_THINKING_BUDGET", Some("10000")),
]);
let mut model_config = ModelConfig::new_or_fail("databricks-claude-3-7-sonnet");
model_config.max_tokens = Some(4096);
model_config = model_config.with_request_params(Some(std::collections::HashMap::from([(
"thinking_type".to_string(),
json!("enabled"),
)])));
let request = create_request(&model_config, "system", &[], &[], &ImageFormat::OpenAi)?;
+4
View File
@@ -149,6 +149,10 @@ pub async fn get_from_registry(name: &str) -> Result<ProviderEntry> {
.cloned()
}
pub async fn inventory_identity(name: &str) -> Result<super::inventory::InventoryIdentityInput> {
get_from_registry(name).await?.inventory_identity()
}
pub async fn create(
name: &str,
model: ModelConfig,
+970
View File
@@ -0,0 +1,970 @@
use super::base::{ConfigKey, ModelInfo};
use super::canonical::{map_provider_name, map_to_canonical_model, CanonicalModelRegistry};
use crate::config::declarative_providers::{DeclarativeProviderConfig, ProviderEngine};
use crate::config::Config;
use crate::session::session_manager::SessionStorage;
use anyhow::Result;
use chrono::{DateTime, Duration, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};
use sqlx::{Pool, Row, Sqlite, Transaction};
use std::collections::{BTreeMap, HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::RwLock;
const STALE_AFTER_HOURS: i64 = 24;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ProviderInventoryEntry {
pub provider_id: String,
pub provider_name: String,
pub configured: bool,
pub supports_refresh: bool,
pub refreshing: bool,
pub models: Vec<InventoryModel>,
pub last_updated_at: Option<DateTime<Utc>>,
pub last_refresh_attempt_at: Option<DateTime<Utc>>,
pub last_refresh_error: Option<String>,
pub model_selection_hint: Option<String>,
}
/// Families whose latest model should be surfaced in the compact picker.
/// Each entry is matched against the `family` field of enriched models.
const RECOMMENDED_FAMILIES: &[&str] = &[
"claude-opus",
"claude-sonnet",
"gpt",
"gpt-mini",
"glm",
"gemini-pro",
"gemini-flash",
"gemma",
];
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InventoryModel {
pub id: String,
pub name: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub family: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub context_limit: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning: Option<bool>,
/// Whether this model should appear in the compact recommended picker.
pub recommended: bool,
}
#[derive(Debug, Clone)]
pub struct InventoryIdentity {
pub provider_id: String,
pub provider_family: String,
pub inventory_key: String,
}
#[derive(Debug, Clone, Default)]
pub struct InventoryIdentityInput {
pub provider_id: String,
pub provider_family: String,
pub public_inputs: BTreeMap<String, String>,
pub secret_inputs: BTreeMap<String, String>,
}
impl InventoryIdentityInput {
pub fn new(
provider_id: impl Into<String>,
provider_family: impl Into<String>,
) -> InventoryIdentityInput {
InventoryIdentityInput {
provider_id: provider_id.into(),
provider_family: provider_family.into(),
public_inputs: BTreeMap::new(),
secret_inputs: BTreeMap::new(),
}
}
pub fn with_public(
mut self,
key: impl Into<String>,
value: impl Into<String>,
) -> InventoryIdentityInput {
self.public_inputs.insert(key.into(), value.into());
self
}
pub fn with_secret(
mut self,
key: impl Into<String>,
value: impl Into<String>,
) -> InventoryIdentityInput {
self.secret_inputs.insert(key.into(), value.into());
self
}
pub fn into_identity(self) -> Result<InventoryIdentity> {
let InventoryIdentityInput {
provider_id,
provider_family,
public_inputs,
secret_inputs,
} = self;
let payload = serde_json::json!({
"provider_family": provider_family,
"public_inputs": public_inputs,
"secret_inputs": secret_inputs,
});
let digest = Sha256::digest(serde_json::to_vec(&payload)?);
Ok(InventoryIdentity {
provider_id,
provider_family,
inventory_key: format!("{digest:x}"),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum RefreshSkipReason {
UnknownProvider,
NotConfigured,
DoesNotSupportRefresh,
AlreadyRefreshing,
}
#[derive(Debug, Clone)]
pub struct RefreshSkip {
pub provider_id: String,
pub reason: RefreshSkipReason,
}
#[derive(Debug, Clone, Default)]
pub struct RefreshPlan {
pub started: Vec<String>,
pub skipped: Vec<RefreshSkip>,
}
#[derive(Clone)]
pub struct ProviderInventoryService {
storage: Arc<SessionStorage>,
refreshing_keys: Arc<RwLock<HashSet<String>>>,
}
#[derive(Debug, Clone)]
struct InventorySnapshot {
models: Vec<InventoryModel>,
last_updated_at: Option<DateTime<Utc>>,
last_refresh_attempt_at: Option<DateTime<Utc>>,
last_refresh_error: Option<String>,
}
#[derive(Debug, Clone)]
struct ProviderDescriptor {
provider_id: String,
provider_name: String,
identity: InventoryIdentity,
configured: bool,
supports_refresh: bool,
static_models: Vec<ModelInfo>,
model_selection_hint: Option<String>,
}
impl ProviderInventoryService {
pub fn new(storage: Arc<SessionStorage>) -> ProviderInventoryService {
ProviderInventoryService {
storage,
refreshing_keys: Arc::new(RwLock::new(HashSet::new())),
}
}
pub async fn entry_for_provider(
&self,
provider_id: &str,
) -> Result<Option<ProviderInventoryEntry>> {
let Some(descriptor) = self.describe_provider(provider_id).await? else {
return Ok(None);
};
let snapshot = self.read_snapshot(&descriptor.identity).await?;
let refreshing = self
.refreshing_keys
.read()
.await
.contains(&descriptor.identity.inventory_key);
let models = inventory_models_from_snapshot(
snapshot.as_ref(),
&descriptor.identity.provider_family,
&descriptor.static_models,
);
Ok(Some(ProviderInventoryEntry {
provider_id: descriptor.provider_id,
provider_name: descriptor.provider_name,
configured: descriptor.configured,
supports_refresh: descriptor.supports_refresh,
refreshing,
models,
last_updated_at: snapshot
.as_ref()
.and_then(|snapshot| snapshot.last_updated_at),
last_refresh_attempt_at: snapshot
.as_ref()
.and_then(|snapshot| snapshot.last_refresh_attempt_at),
last_refresh_error: snapshot.and_then(|snapshot| snapshot.last_refresh_error),
model_selection_hint: descriptor.model_selection_hint,
}))
}
pub async fn entries(&self, provider_ids: &[String]) -> Result<Vec<ProviderInventoryEntry>> {
let ids = self.resolve_provider_ids(provider_ids).await;
let mut entries = Vec::with_capacity(ids.len());
for provider_id in ids {
if let Some(entry) = self.entry_for_provider(&provider_id).await? {
entries.push(entry);
}
}
Ok(entries)
}
pub async fn plan_refresh(&self, provider_ids: &[String]) -> Result<RefreshPlan> {
let ids = self.resolve_provider_ids(provider_ids).await;
let mut plan = RefreshPlan::default();
for provider_id in ids {
let Some(descriptor) = self.describe_provider(&provider_id).await? else {
plan.skipped.push(RefreshSkip {
provider_id,
reason: RefreshSkipReason::UnknownProvider,
});
continue;
};
if !descriptor.supports_refresh {
plan.skipped.push(RefreshSkip {
provider_id: descriptor.provider_id,
reason: RefreshSkipReason::DoesNotSupportRefresh,
});
continue;
}
if !descriptor.configured {
plan.skipped.push(RefreshSkip {
provider_id: descriptor.provider_id,
reason: RefreshSkipReason::NotConfigured,
});
continue;
}
let mut refreshing_keys = self.refreshing_keys.write().await;
if refreshing_keys.contains(&descriptor.identity.inventory_key) {
plan.skipped.push(RefreshSkip {
provider_id: descriptor.provider_id,
reason: RefreshSkipReason::AlreadyRefreshing,
});
continue;
}
refreshing_keys.insert(descriptor.identity.inventory_key.clone());
drop(refreshing_keys);
self.mark_refresh_started(&descriptor.identity).await?;
plan.started.push(descriptor.provider_id);
}
Ok(plan)
}
pub async fn store_refreshed_models(
&self,
provider_id: &str,
model_ids: &[String],
) -> Result<()> {
let descriptor = self.require_provider(provider_id).await?;
let models =
enrich_model_ids_with_canonical(&descriptor.identity.provider_family, model_ids);
let now = Utc::now();
let pool = self.storage.pool().await?;
let mut tx = pool.begin().await?;
sqlx::query(
r#"
INSERT INTO provider_inventory_entries (
inventory_key,
provider_id,
provider_family,
last_updated_at,
last_refresh_attempt_at,
last_refresh_error,
updated_at
) VALUES (?, ?, ?, ?, ?, NULL, CURRENT_TIMESTAMP)
ON CONFLICT(inventory_key) DO UPDATE SET
provider_id = excluded.provider_id,
provider_family = excluded.provider_family,
last_updated_at = excluded.last_updated_at,
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
last_refresh_error = NULL,
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&descriptor.identity.inventory_key)
.bind(&descriptor.identity.provider_id)
.bind(&descriptor.identity.provider_family)
.bind(now.to_rfc3339())
.bind(now.to_rfc3339())
.execute(&mut *tx)
.await?;
sqlx::query("DELETE FROM provider_inventory_models WHERE inventory_key = ?")
.bind(&descriptor.identity.inventory_key)
.execute(&mut *tx)
.await?;
for (ordinal, model) in models.iter().enumerate() {
sqlx::query(
r#"
INSERT INTO provider_inventory_models (
inventory_key,
ordinal,
model_id,
name,
family,
context_limit,
reasoning,
recommended
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)
"#,
)
.bind(&descriptor.identity.inventory_key)
.bind(i64::try_from(ordinal)?)
.bind(&model.id)
.bind(&model.name)
.bind(&model.family)
.bind(model.context_limit.map(i64::try_from).transpose()?)
.bind(model.reasoning)
.bind(model.recommended)
.execute(&mut *tx)
.await?;
}
tx.commit().await?;
self.refreshing_keys
.write()
.await
.remove(&descriptor.identity.inventory_key);
Ok(())
}
pub async fn store_refresh_error(
&self,
provider_id: &str,
error: impl Into<String>,
) -> Result<()> {
let descriptor = self.require_provider(provider_id).await?;
let error = error.into();
let existing = self.read_snapshot(&descriptor.identity).await?;
sqlx::query(
r#"
INSERT INTO provider_inventory_entries (
inventory_key,
provider_id,
provider_family,
last_updated_at,
last_refresh_attempt_at,
last_refresh_error,
updated_at
) VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
ON CONFLICT(inventory_key) DO UPDATE SET
provider_id = excluded.provider_id,
provider_family = excluded.provider_family,
last_updated_at = excluded.last_updated_at,
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
last_refresh_error = excluded.last_refresh_error,
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&descriptor.identity.inventory_key)
.bind(&descriptor.identity.provider_id)
.bind(&descriptor.identity.provider_family)
.bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339())))
.bind(Utc::now().to_rfc3339())
.bind(error)
.execute(self.storage.pool().await?)
.await?;
self.refreshing_keys
.write()
.await
.remove(&descriptor.identity.inventory_key);
Ok(())
}
pub fn is_stale(entry: &ProviderInventoryEntry) -> bool {
let Some(last_updated_at) = entry.last_updated_at else {
return false;
};
entry.supports_refresh && Utc::now() - last_updated_at > Duration::hours(STALE_AFTER_HOURS)
}
async fn describe_provider(&self, provider_id: &str) -> Result<Option<ProviderDescriptor>> {
let entry = match crate::providers::get_from_registry(provider_id).await {
Ok(entry) => entry,
Err(_) => return Ok(None),
};
let metadata = entry.metadata().clone();
let identity = crate::providers::inventory_identity(provider_id)
.await
.unwrap_or_else(|_| fallback_inventory_identity(provider_id))
.into_identity()?;
Ok(Some(ProviderDescriptor {
provider_id: metadata.name.clone(),
provider_name: metadata.display_name.clone(),
identity,
configured: entry.inventory_configured(),
supports_refresh: entry.supports_inventory_refresh(),
static_models: metadata.known_models,
model_selection_hint: metadata.model_selection_hint,
}))
}
async fn require_provider(&self, provider_id: &str) -> Result<ProviderDescriptor> {
self.describe_provider(provider_id)
.await?
.ok_or_else(|| anyhow::anyhow!("Unknown provider: {}", provider_id))
}
async fn mark_refresh_started(&self, identity: &InventoryIdentity) -> Result<()> {
let existing = self.read_snapshot(identity).await?;
sqlx::query(
r#"
INSERT INTO provider_inventory_entries (
inventory_key,
provider_id,
provider_family,
last_updated_at,
last_refresh_attempt_at,
last_refresh_error,
updated_at
) VALUES (?, ?, ?, ?, ?, NULL, CURRENT_TIMESTAMP)
ON CONFLICT(inventory_key) DO UPDATE SET
provider_id = excluded.provider_id,
provider_family = excluded.provider_family,
last_updated_at = excluded.last_updated_at,
last_refresh_attempt_at = excluded.last_refresh_attempt_at,
last_refresh_error = NULL,
updated_at = CURRENT_TIMESTAMP
"#,
)
.bind(&identity.inventory_key)
.bind(&identity.provider_id)
.bind(&identity.provider_family)
.bind(existing.and_then(|snapshot| snapshot.last_updated_at.map(|time| time.to_rfc3339())))
.bind(Utc::now().to_rfc3339())
.execute(self.storage.pool().await?)
.await?;
Ok(())
}
async fn read_snapshot(
&self,
identity: &InventoryIdentity,
) -> Result<Option<InventorySnapshot>> {
let pool = self.storage.pool().await?;
let entry = sqlx::query(
r#"
SELECT last_updated_at, last_refresh_attempt_at, last_refresh_error
FROM provider_inventory_entries
WHERE inventory_key = ?
"#,
)
.bind(&identity.inventory_key)
.fetch_optional(pool)
.await?;
let Some(entry) = entry else {
return Ok(None);
};
let last_updated_at = parse_optional_datetime(entry.try_get("last_updated_at")?)?;
let last_refresh_attempt_at =
parse_optional_datetime(entry.try_get("last_refresh_attempt_at")?)?;
let last_refresh_error = entry.try_get("last_refresh_error")?;
let rows = sqlx::query(
r#"
SELECT model_id, name, family, context_limit, reasoning, recommended
FROM provider_inventory_models
WHERE inventory_key = ?
ORDER BY ordinal
"#,
)
.bind(&identity.inventory_key)
.fetch_all(pool)
.await?;
let models = rows
.into_iter()
.map(|row| {
Ok(InventoryModel {
id: row.try_get("model_id")?,
name: row.try_get("name")?,
family: row.try_get("family")?,
context_limit: row
.try_get::<Option<i64>, _>("context_limit")?
.map(usize::try_from)
.transpose()?,
reasoning: row.try_get("reasoning")?,
recommended: row
.try_get::<Option<bool>, _>("recommended")?
.unwrap_or(false),
})
})
.collect::<Result<Vec<_>, anyhow::Error>>()?;
Ok(Some(InventorySnapshot {
models,
last_updated_at,
last_refresh_attempt_at,
last_refresh_error,
}))
}
async fn resolve_provider_ids(&self, provider_ids: &[String]) -> Vec<String> {
let mut ids = if provider_ids.is_empty() {
crate::providers::providers()
.await
.into_iter()
.map(|(metadata, _)| metadata.name)
.collect::<Vec<_>>()
} else {
provider_ids.to_vec()
};
ids.sort();
ids.dedup();
ids
}
}
pub fn default_inventory_identity(
provider_id: &str,
provider_family: &str,
config_keys: &[ConfigKey],
config: &Config,
) -> InventoryIdentityInput {
let mut identity = InventoryIdentityInput::new(provider_id, provider_family);
for key in config_keys {
if key.secret {
if let Some(value) = config_secret_value(config, &key.name) {
identity.secret_inputs.insert(key.name.clone(), value);
}
} else if let Some(value) = config_param_value(config, &key.name) {
identity.public_inputs.insert(key.name.clone(), value);
}
}
identity
}
pub fn default_inventory_configured(config_keys: &[ConfigKey], config: &Config) -> bool {
config_keys.iter().all(|key| {
if !key.required {
return true;
}
if key.default.is_some() {
return true;
}
if key.secret {
config.get_secret::<serde_json::Value>(&key.name).is_ok()
} else {
config.get_param::<serde_json::Value>(&key.name).is_ok()
}
})
}
pub fn declarative_inventory_identity(
config: &DeclarativeProviderConfig,
) -> Result<InventoryIdentityInput> {
let global = Config::global();
let mut identity = InventoryIdentityInput::new(
config.name.clone(),
config
.catalog_provider_id
.clone()
.unwrap_or_else(|| match config.engine {
ProviderEngine::OpenAI => "openai".to_string(),
ProviderEngine::Anthropic => "anthropic".to_string(),
ProviderEngine::Ollama => "ollama".to_string(),
}),
);
identity
.public_inputs
.insert("base_url".to_string(), config.base_url.clone());
if let Some(base_path) = &config.base_path {
identity
.public_inputs
.insert("base_path".to_string(), base_path.clone());
}
if let Some(catalog_provider_id) = &config.catalog_provider_id {
identity.public_inputs.insert(
"catalog_provider_id".to_string(),
catalog_provider_id.clone(),
);
}
if let Some(dynamic_models) = config.dynamic_models {
identity
.public_inputs
.insert("dynamic_models".to_string(), dynamic_models.to_string());
}
identity.public_inputs.insert(
"skip_canonical_filtering".to_string(),
config.skip_canonical_filtering.to_string(),
);
if !config.models.is_empty() {
identity.public_inputs.insert(
"models".to_string(),
serde_json::to_string(
&config
.models
.iter()
.map(|model| &model.name)
.collect::<Vec<_>>(),
)?,
);
}
if let Some(headers) = &config.headers {
identity
.public_inputs
.insert("headers".to_string(), serialize_string_map(headers)?);
}
if config.requires_auth && !config.api_key_env.is_empty() {
if let Some(value) = config_secret_value(global, &config.api_key_env) {
identity
.secret_inputs
.insert(config.api_key_env.clone(), value);
}
}
Ok(identity)
}
pub fn config_param_value(config: &Config, key: &str) -> Option<String> {
config
.get_param::<serde_json::Value>(key)
.ok()
.and_then(|value| normalize_json_value(&value))
}
pub fn config_secret_value(config: &Config, key: &str) -> Option<String> {
config
.get_secret::<serde_json::Value>(key)
.ok()
.and_then(|value| normalize_json_value(&value))
}
pub fn serialize_string_map(map: &HashMap<String, String>) -> Result<String> {
let ordered = map
.iter()
.map(|(key, value)| (key.clone(), value.clone()))
.collect::<BTreeMap<_, _>>();
Ok(serde_json::to_string(&ordered)?)
}
fn parse_optional_datetime(value: Option<String>) -> Result<Option<DateTime<Utc>>> {
value
.map(|value| value.parse::<DateTime<Utc>>())
.transpose()
.map_err(Into::into)
}
fn normalize_json_value(value: &serde_json::Value) -> Option<String> {
match value {
serde_json::Value::Null => None,
serde_json::Value::String(value) if value.is_empty() => None,
serde_json::Value::String(value) => Some(value.clone()),
other => serde_json::to_string(other).ok(),
}
}
fn fallback_inventory_identity(provider_id: &str) -> InventoryIdentityInput {
InventoryIdentityInput::new(
provider_id.to_string(),
map_provider_name(provider_id).to_string(),
)
}
fn enrich_model_ids_with_canonical(
provider_family: &str,
model_ids: &[String],
) -> Vec<InventoryModel> {
let mut models: Vec<InventoryModel> = Vec::new();
let mut seen_names: HashSet<String> = HashSet::new();
for id in model_ids {
let model = enriched_model(provider_family, id, None);
if !seen_names.insert(model.name.clone()) {
continue;
}
models.push(model);
}
// For databricks, prefer goose- prefixed model_ids when there are duplicates.
// Re-scan: if a later model_id with "goose-" prefix maps to the same display name,
// swap it in.
if provider_family == "databricks" {
let mut name_to_idx: HashMap<String, usize> = HashMap::new();
for (idx, model) in models.iter().enumerate() {
name_to_idx.insert(model.name.clone(), idx);
}
for id in model_ids {
if !id.starts_with("goose-") {
continue;
}
let candidate = enriched_model(provider_family, id, None);
if let Some(&idx) = name_to_idx.get(&candidate.name) {
if !models[idx].id.starts_with("goose-") {
models[idx].id = candidate.id;
}
}
}
}
// Mark the latest model per recommended family.
let mut seen_recommended_families: HashSet<String> = HashSet::new();
for model in &mut models {
if let Some(family) = &model.family {
if RECOMMENDED_FAMILIES.contains(&family.as_str())
&& seen_recommended_families.insert(family.clone())
{
model.recommended = true;
}
}
}
models
}
fn configured_models_to_inventory(
provider_family: &str,
models: &[ModelInfo],
) -> Vec<InventoryModel> {
let mut result: Vec<InventoryModel> = Vec::new();
let mut seen_names: HashSet<String> = HashSet::new();
for model in models {
let enriched = enriched_model(provider_family, &model.name, Some(model.context_limit));
if seen_names.insert(enriched.name.clone()) {
result.push(enriched);
}
}
let mut seen_recommended_families: HashSet<String> = HashSet::new();
for model in &mut result {
if let Some(family) = &model.family {
if RECOMMENDED_FAMILIES.contains(&family.as_str())
&& seen_recommended_families.insert(family.clone())
{
model.recommended = true;
}
}
}
result
}
fn inventory_models_from_snapshot(
snapshot: Option<&InventorySnapshot>,
provider_family: &str,
configured_models: &[ModelInfo],
) -> Vec<InventoryModel> {
match snapshot {
Some(snapshot) if !snapshot.models.is_empty() || snapshot.last_updated_at.is_some() => {
snapshot.models.clone()
}
_ => configured_models_to_inventory(provider_family, configured_models),
}
}
fn enriched_model(
provider_family: &str,
model_id: &str,
fallback_context_limit: Option<usize>,
) -> InventoryModel {
let registry = CanonicalModelRegistry::bundled().ok();
let canonical = registry.as_ref().and_then(|registry| {
let canonical_id = map_to_canonical_model(provider_family, model_id, registry)?;
let (provider, model) = canonical_id.split_once('/')?;
registry.get(provider, model).cloned()
});
InventoryModel {
id: model_id.to_string(),
name: canonical
.as_ref()
.map(|model| model.name.clone())
.unwrap_or_else(|| model_id.to_string()),
family: canonical.as_ref().and_then(|model| model.family.clone()),
context_limit: canonical
.as_ref()
.map(|model| model.limit.context)
.or(fallback_context_limit),
reasoning: canonical.as_ref().and_then(|model| model.reasoning),
recommended: false,
}
}
pub async fn create_tables(pool: &Pool<Sqlite>) -> Result<()> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS provider_inventory_entries (
inventory_key TEXT PRIMARY KEY,
provider_id TEXT NOT NULL,
provider_family TEXT NOT NULL,
last_updated_at TEXT,
last_refresh_attempt_at TEXT,
last_refresh_error TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"#,
)
.execute(pool)
.await?;
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS provider_inventory_models (
inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE,
ordinal INTEGER NOT NULL,
model_id TEXT NOT NULL,
name TEXT NOT NULL,
family TEXT,
context_limit INTEGER,
reasoning BOOLEAN,
recommended BOOLEAN,
PRIMARY KEY (inventory_key, ordinal)
)
"#,
)
.execute(pool)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)",
)
.execute(pool)
.await?;
Ok(())
}
pub async fn create_tables_in_tx(tx: &mut Transaction<'_, Sqlite>) -> Result<()> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS provider_inventory_entries (
inventory_key TEXT PRIMARY KEY,
provider_id TEXT NOT NULL,
provider_family TEXT NOT NULL,
last_updated_at TEXT,
last_refresh_attempt_at TEXT,
last_refresh_error TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
)
"#,
)
.execute(&mut **tx)
.await?;
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS provider_inventory_models (
inventory_key TEXT NOT NULL REFERENCES provider_inventory_entries(inventory_key) ON DELETE CASCADE,
ordinal INTEGER NOT NULL,
model_id TEXT NOT NULL,
name TEXT NOT NULL,
family TEXT,
context_limit INTEGER,
reasoning BOOLEAN,
recommended BOOLEAN,
PRIMARY KEY (inventory_key, ordinal)
)
"#,
)
.execute(&mut **tx)
.await?;
sqlx::query(
"CREATE INDEX IF NOT EXISTS idx_provider_inventory_provider_id ON provider_inventory_entries(provider_id)",
)
.execute(&mut **tx)
.await?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn inventory_identity_hash_changes_with_secret_inputs() {
let left = InventoryIdentityInput::new("openai", "openai")
.with_public("host", "https://api.openai.com")
.with_secret("api_key", "secret-a")
.into_identity()
.unwrap();
let right = InventoryIdentityInput::new("openai", "openai")
.with_public("host", "https://api.openai.com")
.with_secret("api_key", "secret-b")
.into_identity()
.unwrap();
assert_ne!(left.inventory_key, right.inventory_key);
}
#[test]
fn configured_models_use_canonical_enrichment() {
let models =
configured_models_to_inventory("anthropic", &[ModelInfo::new("claude-sonnet-4-5", 0)]);
assert_eq!(models.len(), 1);
assert!(models[0].name.contains("Claude"));
}
#[test]
fn inventory_uses_configured_models_before_first_successful_refresh() {
let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)];
let snapshot = InventorySnapshot {
models: vec![],
last_updated_at: None,
last_refresh_attempt_at: Some(Utc::now()),
last_refresh_error: Some("auth failed".to_string()),
};
let models =
inventory_models_from_snapshot(Some(&snapshot), "anthropic", &configured_models);
assert_eq!(models.len(), 1);
assert_eq!(models[0].id, "claude-sonnet-4-5");
}
#[test]
fn inventory_preserves_empty_models_after_successful_refresh() {
let configured_models = [ModelInfo::new("claude-sonnet-4-5", 0)];
let snapshot = InventorySnapshot {
models: vec![],
last_updated_at: Some(Utc::now()),
last_refresh_attempt_at: Some(Utc::now()),
last_refresh_error: None,
};
let models =
inventory_models_from_snapshot(Some(&snapshot), "anthropic", &configured_models);
assert!(models.is_empty());
}
}
+3 -1
View File
@@ -1,3 +1,4 @@
mod acp_tooling;
pub mod amp_acp;
pub mod anthropic;
pub mod api_client;
@@ -28,6 +29,7 @@ pub mod gemini_oauth;
pub mod githubcopilot;
pub mod google;
mod init;
pub mod inventory;
pub mod kimicode;
pub mod litellm;
#[cfg(feature = "local-inference")]
@@ -56,6 +58,6 @@ pub mod xai;
pub use init::{
cleanup_provider, create, create_with_default_model, create_with_named_model,
get_from_registry, providers, refresh_custom_providers,
get_from_registry, inventory_identity, providers, refresh_custom_providers,
};
pub use retry::{retry_operation, RetryConfig};
+17
View File
@@ -1,6 +1,7 @@
use super::api_client::{ApiClient, AuthMethod};
use super::base::{ConfigKey, MessageStream, Provider, ProviderDef, ProviderMetadata};
use super::errors::ProviderError;
use super::inventory::InventoryIdentityInput;
use super::openai_compatible::handle_status_openai_compat;
use super::retry::{ProviderRetry, RetryConfig};
use super::utils::{ImageFormat, RequestLog};
@@ -256,6 +257,22 @@ impl ProviderDef for OllamaProvider {
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(Self::from_env(model))
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
let config = crate::config::Config::global();
Ok(
InventoryIdentityInput::new(OLLAMA_PROVIDER_NAME, OLLAMA_PROVIDER_NAME).with_public(
"host",
config
.get_param::<String>("OLLAMA_HOST")
.unwrap_or_else(|_| OLLAMA_HOST.to_string()),
),
)
}
}
#[async_trait]
+53
View File
@@ -7,6 +7,7 @@ use super::formats::openai_responses::{
create_responses_request, get_responses_usage, responses_api_to_message,
responses_api_to_streaming_message, ResponsesApiResponse,
};
use super::inventory::{config_secret_value, InventoryIdentityInput};
use super::openai_compatible::{
handle_response_openai_compat, handle_status_openai_compat, stream_openai_compat,
};
@@ -425,6 +426,58 @@ impl ProviderDef for OpenAiProvider {
) -> BoxFuture<'static, Result<Self::Provider>> {
Box::pin(Self::from_env(model))
}
fn supports_inventory_refresh() -> bool {
true
}
fn inventory_configured() -> bool {
let config = crate::config::Config::global();
// If the host is explicitly set to something non-default, trust the user's
// custom setup (e.g. a local server that doesn't require an API key).
if let Ok(host) = config.get_param::<String>("OPENAI_HOST") {
if host != "https://api.openai.com" {
return true;
}
}
// Standard OpenAI endpoint requires an API key.
config
.get_secret::<serde_json::Value>("OPENAI_API_KEY")
.is_ok()
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
let config = crate::config::Config::global();
let mut identity =
InventoryIdentityInput::new(OPEN_AI_PROVIDER_NAME, OPEN_AI_PROVIDER_NAME)
.with_public(
"host",
config
.get_param::<String>("OPENAI_HOST")
.unwrap_or_else(|_| "https://api.openai.com".to_string()),
)
.with_public(
"base_path",
config
.get_param::<String>("OPENAI_BASE_PATH")
.unwrap_or_else(|_| OPEN_AI_DEFAULT_BASE_PATH.to_string()),
);
if let Ok(organization) = config.get_param::<String>("OPENAI_ORGANIZATION") {
identity = identity.with_public("organization", organization);
}
if let Ok(project) = config.get_param::<String>("OPENAI_PROJECT") {
identity = identity.with_public("project", project);
}
if let Some(api_key) = config_secret_value(config, "OPENAI_API_KEY") {
identity = identity.with_secret("api_key", api_key);
}
if let Some(custom_headers) = config_secret_value(config, "OPENAI_CUSTOM_HEADERS") {
identity = identity.with_secret("custom_headers", custom_headers);
}
Ok(identity)
}
}
#[async_trait]
+15
View File
@@ -9,7 +9,9 @@ use crate::acp::{
use crate::config::search_path::SearchPaths;
use crate::config::{Config, GooseMode};
use crate::model::ModelConfig;
use crate::providers::acp_tooling::{acp_adapter_installed, acp_inventory_identity};
use crate::providers::base::{ProviderDef, ProviderMetadata};
use crate::providers::inventory::InventoryIdentityInput;
const PI_ACP_PROVIDER_NAME: &str = "pi-acp";
const PI_ACP_DOC_URL: &str = "https://github.com/anthropics/pi";
@@ -36,6 +38,7 @@ impl ProviderDef for PiAcpProvider {
"Set in your goose config file (`~/.config/goose/config.yaml` on macOS/Linux):\n GOOSE_PROVIDER: pi-acp\n GOOSE_MODEL: current",
"Restart goose for changes to take effect",
])
.with_model_selection_hint("Use the Pi CLI to configure models")
}
fn from_env(
@@ -70,4 +73,16 @@ impl ProviderDef for PiAcpProvider {
AcpProvider::connect(metadata.name, model, goose_mode, provider_config).await
})
}
fn supports_inventory_refresh() -> bool {
false
}
fn inventory_identity() -> Result<InventoryIdentityInput> {
acp_inventory_identity(PI_ACP_PROVIDER_NAME, PI_ACP_BINARY)
}
fn inventory_configured() -> bool {
acp_adapter_installed(PI_ACP_BINARY)
}
}
@@ -1,4 +1,5 @@
use super::base::{ModelInfo, Provider, ProviderDef, ProviderMetadata, ProviderType};
use super::inventory::InventoryIdentityInput;
use crate::config::{DeclarativeProviderConfig, ExtensionConfig};
use crate::model::ModelConfig;
use anyhow::Result;
@@ -14,12 +15,20 @@ pub type ProviderConstructor = Arc<
pub type ProviderCleanup = Arc<dyn Fn() -> BoxFuture<'static, Result<()>> + Send + Sync>;
pub type ProviderInventoryIdentityResolver =
Arc<dyn Fn() -> Result<InventoryIdentityInput> + Send + Sync>;
pub type ProviderInventoryConfiguredResolver = Arc<dyn Fn() -> bool + Send + Sync>;
#[derive(Clone)]
pub struct ProviderEntry {
metadata: ProviderMetadata,
pub(crate) constructor: ProviderConstructor,
pub(crate) inventory_identity: ProviderInventoryIdentityResolver,
pub(crate) inventory_configured: ProviderInventoryConfiguredResolver,
pub(crate) cleanup: Option<ProviderCleanup>,
provider_type: ProviderType,
supports_inventory_refresh: bool,
}
impl ProviderEntry {
@@ -27,6 +36,22 @@ impl ProviderEntry {
&self.metadata
}
pub fn provider_type(&self) -> ProviderType {
self.provider_type
}
pub fn supports_inventory_refresh(&self) -> bool {
self.supports_inventory_refresh
}
pub fn inventory_identity(&self) -> Result<InventoryIdentityInput> {
(self.inventory_identity)()
}
pub fn inventory_configured(&self) -> bool {
(self.inventory_configured)()
}
fn normalize_model_config(&self, mut model: ModelConfig) -> ModelConfig {
model = model.with_canonical_limits(&self.metadata.name);
@@ -92,24 +117,30 @@ impl ProviderRegistry {
Ok(Arc::new(provider) as Arc<dyn Provider>)
})
}),
inventory_identity: Arc::new(F::inventory_identity),
inventory_configured: Arc::new(F::inventory_configured),
cleanup: None,
provider_type: if preferred {
ProviderType::Preferred
} else {
ProviderType::Builtin
},
supports_inventory_refresh: F::supports_inventory_refresh(),
},
);
}
pub fn register_with_name<P, F>(
pub fn register_with_name<P, F, G>(
&mut self,
config: &DeclarativeProviderConfig,
provider_type: ProviderType,
supports_inventory_refresh: bool,
constructor: F,
inventory_identity: G,
) 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
@@ -174,7 +205,9 @@ impl ProviderRegistry {
model_doc_link: base_metadata.model_doc_link,
config_keys,
setup_steps: vec![],
model_selection_hint: None,
};
let inventory_config_keys = custom_metadata.config_keys.clone();
self.entries.insert(
config.name.clone(),
@@ -187,8 +220,16 @@ impl ProviderRegistry {
Ok(Arc::new(provider) as Arc<dyn Provider>)
})
}),
inventory_identity: Arc::new(inventory_identity),
inventory_configured: Arc::new(move || {
super::inventory::default_inventory_configured(
&inventory_config_keys,
crate::config::Config::global(),
)
}),
cleanup: None,
provider_type,
supports_inventory_refresh,
},
);
}
+6 -1
View File
@@ -19,7 +19,7 @@ use std::sync::{Arc, LazyLock};
use tracing::{info, warn};
use utoipa::ToSchema;
pub const CURRENT_SCHEMA_VERSION: i32 = 10;
pub const CURRENT_SCHEMA_VERSION: i32 = 11;
pub const SESSIONS_FOLDER: &str = "sessions";
pub const DB_NAME: &str = "sessions.db";
@@ -717,6 +717,8 @@ impl SessionStorage {
.execute(pool)
.await?;
crate::providers::inventory::create_tables(pool).await?;
Ok(())
}
@@ -1060,6 +1062,9 @@ impl SessionStorage {
.execute(&mut **tx)
.await?;
}
11 => {
crate::providers::inventory::create_tables_in_tx(tx).await?;
}
_ => {
anyhow::bail!("Unknown migration version: {}", version);
}