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

feat: extensions read config (#1637)

Browse Source
This commit is contained in:
Lily Delalande
2025-03-12 18:41:31 -07:00
committed by GitHub
parent 68b3b3f9cc
commit 4537264387
15 changed files with 756 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-cli/src/commands/configure.rs
+14 -3
View File
@@ -1,6 +1,7 @@
use cliclack::spinner;
use console::style;
use goose::agents::{extension::Envs, ExtensionConfig};
use goose::config::extensions::name_to_key;
use goose::config::{Config, ConfigError, ExperimentManager, ExtensionEntry, ExtensionManager};
use goose::message::Message;
use goose::providers::{create, providers};
@@ -387,7 +388,10 @@ pub fn toggle_extensions_dialog() -> Result<(), Box<dyn Error>> {
// Update enabled status for each extension
for name in extension_status.iter().map(|(name, _)| name) {
ExtensionManager::set_enabled(name, selected.iter().any(|s| s.as_str() == name))?;
ExtensionManager::set_enabled(
&name_to_key(name),
selected.iter().any(|s| s.as_str() == name),
)?;
}
cliclack::outro("Extension settings updated successfully")?;
@@ -638,10 +642,17 @@ pub fn remove_extension_dialog() -> Result<(), Box<dyn Error>> {
return Ok(());
}
// Filter out only disabled extensions
let disabled_extensions: Vec<_> = extensions
.iter()
.filter(|entry| !entry.enabled)
.map(|entry| (entry.config.name().to_string(), entry.enabled))
.collect();
let selected = cliclack::multiselect("Select extensions to remove (note: you can only remove disabled extensions - use \"space\" to toggle and \"enter\" to submit)")
.required(false)
.items(
&extension_status
&disabled_extensions
.iter()
.filter(|(_, enabled)| !enabled)
.map(|(name, _)| (name, name.as_str(), ""))
@@ -650,7 +661,7 @@ pub fn remove_extension_dialog() -> Result<(), Box<dyn Error>> {
.interact()?;
for name in selected {
ExtensionManager::remove(name)?;
ExtensionManager::remove(&name_to_key(name))?;
cliclack::outro(format!("Removed {} extension", style(name).green()))?;
}
crates/goose-server/src/openapi.rs
+11 -2
View File
@@ -1,5 +1,8 @@
use utoipa::OpenApi;
use goose::agents::extension::Envs;
use goose::agents::ExtensionConfig;
use goose::config::ExtensionEntry;
use goose::providers::base::ConfigKey;
use goose::providers::base::ProviderMetadata;
@@ -12,6 +15,8 @@ use goose::providers::base::ProviderMetadata;
super::routes::config_management::read_config,
super::routes::config_management::add_extension,
super::routes::config_management::remove_extension,
super::routes::config_management::toggle_extension,
super::routes::config_management::get_extensions,
super::routes::config_management::update_extension,
super::routes::config_management::read_all_config,
super::routes::config_management::providers
@@ -19,13 +24,17 @@ use goose::providers::base::ProviderMetadata;
components(schemas(
super::routes::config_management::UpsertConfigQuery,
super::routes::config_management::ConfigKeyQuery,
super::routes::config_management::ExtensionQuery,
super::routes::config_management::ConfigResponse,
super::routes::config_management::ProvidersResponse,
super::routes::config_management::ProvidersResponse,
super::routes::config_management::ProviderDetails,
super::routes::config_management::ExtensionResponse,
super::routes::config_management::ExtensionQuery,
ProviderMetadata,
ConfigKey
ExtensionEntry,
ExtensionConfig,
ConfigKey,
Envs,
))
)]
pub struct ApiDoc;
crates/goose-server/src/routes/config_management.rs
+136 -95
View File
@@ -1,10 +1,15 @@
use crate::routes::utils::check_provider_configured;
use crate::state::AppState;
use axum::routing::put;
use axum::{
extract::State,
routing::{delete, get, post},
Json, Router,
};
use goose::agents::ExtensionConfig;
use goose::config::extensions::name_to_key;
use goose::config::Config;
use goose::config::{ExtensionEntry, ExtensionManager};
use goose::providers::base::ProviderMetadata;
use goose::providers::providers as get_providers;
use http::{HeaderMap, StatusCode};
@@ -13,9 +18,6 @@ use serde_json::Value;
use std::collections::HashMap;
use utoipa::ToSchema;
use crate::routes::utils::check_provider_configured;
use crate::state::AppState;
fn verify_secret_key(headers: &HeaderMap, state: &AppState) -> Result<StatusCode, StatusCode> {
// Verify secret key
let secret_key = headers
@@ -30,6 +32,18 @@ fn verify_secret_key(headers: &HeaderMap, state: &AppState) -> Result<StatusCode
}
}
#[derive(Serialize, ToSchema)]
pub struct ExtensionResponse {
pub extensions: Vec<ExtensionEntry>,
}
#[derive(Deserialize, ToSchema)]
pub struct ExtensionQuery {
pub name: String,
pub config: ExtensionConfig,
pub enabled: bool,
}
#[derive(Deserialize, ToSchema)]
pub struct UpsertConfigQuery {
pub key: String,
@@ -43,12 +57,6 @@ pub struct ConfigKeyQuery {
pub is_secret: bool,
}
#[derive(Deserialize, ToSchema)]
pub struct ExtensionQuery {
pub name: String,
pub config: Value,
}
#[derive(Serialize, ToSchema)]
pub struct ConfigResponse {
pub config: HashMap<String, Value>,
@@ -155,9 +163,29 @@ pub async fn read_config(
}
}
#[utoipa::path(
get,
path = "/config/extensions",
responses(
(status = 200, description = "All extensions retrieved successfully", body = ExtensionResponse),
(status = 500, description = "Internal server error")
)
)]
pub async fn get_extensions(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<ExtensionResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
match ExtensionManager::get_all() {
Ok(extensions) => Ok(Json(ExtensionResponse { extensions })),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
#[utoipa::path(
post,
path = "/config/extension",
path = "/config/extensions",
request_body = ExtensionQuery,
responses(
(status = 200, description = "Extension added successfully", body = String),
@@ -168,35 +196,23 @@ pub async fn read_config(
pub async fn add_extension(
State(state): State<AppState>,
headers: HeaderMap,
Json(extension): Json<ExtensionQuery>,
Json(extension_query): Json<ExtensionQuery>,
) -> Result<Json<String>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
// Get current extensions or initialize empty map
let mut extensions: HashMap<String, Value> = config
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
// Add new extension
extensions.insert(extension.name.clone(), extension.config);
// Save updated extensions
match config.set_param(
"extensions",
Value::Object(extensions.into_iter().collect()),
) {
Ok(_) => Ok(Json(format!("Added extension {}", extension.name))),
// Use ExtensionManager to set the extension
match ExtensionManager::set(ExtensionEntry {
enabled: extension_query.enabled,
config: extension_query.config,
}) {
Ok(_) => Ok(Json(format!("Added extension {}", extension_query.name))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
#[utoipa::path(
delete,
path = "/config/extension",
request_body = ConfigKeyQuery,
path = "/config/extensions/{name}",
responses(
(status = 200, description = "Extension removed successfully", body = String),
(status = 404, description = "Extension not found"),
@@ -206,31 +222,98 @@ pub async fn add_extension(
pub async fn remove_extension(
State(state): State<AppState>,
headers: HeaderMap,
Json(query): Json<ConfigKeyQuery>,
axum::extract::Path(name): axum::extract::Path<String>,
) -> Result<Json<String>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
let key = name_to_key(&name);
// Use ExtensionManager to remove the extension
match ExtensionManager::remove(&key) {
Ok(_) => Ok(Json(format!("Removed extension {}", name))),
Err(_) => Err(StatusCode::NOT_FOUND),
}
}
// Get current extensions
let mut extensions: HashMap<String, Value> = match config.get_param("extensions") {
Ok(exts) => exts,
Err(_) => return Err(StatusCode::NOT_FOUND),
#[utoipa::path(
put,
path = "/config/extensions/{name}",
request_body = ExtensionQuery,
responses(
(status = 200, description = "Extension updated successfully", body = String),
(status = 404, description = "Extension not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn update_extension(
State(state): State<AppState>,
headers: HeaderMap,
axum::extract::Path(name): axum::extract::Path<String>,
Json(extension_query): Json<ExtensionQuery>,
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let key = name_to_key(&name);
// Check if extension exists
let extensions = ExtensionManager::get_all().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
if !extensions.iter().any(|entry| entry.config.key() == key) {
return Err(StatusCode::NOT_FOUND);
}
// Use ExtensionManager to update the extension
match ExtensionManager::set(ExtensionEntry {
enabled: extension_query.enabled,
config: extension_query.config,
}) {
Ok(_) => Ok(Json(format!("Updated extension {}", extension_query.name))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
#[utoipa::path(
post,
path = "/extensions/{name}/toggle",
responses(
(status = 200, description = "Extension toggled successfully", body = String),
(status = 404, description = "Extension not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn toggle_extension(
State(state): State<AppState>,
headers: HeaderMap,
axum::extract::Path(name): axum::extract::Path<String>,
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let key = name_to_key(&name);
// Get the extension
let extensions = ExtensionManager::get_all().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
let extension = extensions
.iter()
.find(|e| e.config.key() == key)
.ok_or(StatusCode::NOT_FOUND)?;
// Create a new entry with toggled enabled state
let updated_entry = ExtensionEntry {
enabled: !extension.enabled,
config: extension.config.clone(),
};
// Remove extension if it exists
if extensions.remove(&query.key).is_some() {
// Save updated extensions
match config.set_param(
"extensions",
Value::Object(extensions.into_iter().collect()),
) {
Ok(_) => Ok(Json(format!("Removed extension {}", query.key))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
// Update using ExtensionManager
match ExtensionManager::set(updated_entry) {
Ok(_) => {
let status = if !extension.enabled {
"enabled"
} else {
"disabled"
};
Ok(Json(format!("Extension {} {}", name, status)))
}
} else {
Err(StatusCode::NOT_FOUND)
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
@@ -256,50 +339,6 @@ pub async fn read_all_config(
Ok(Json(ConfigResponse { config: values }))
}
#[utoipa::path(
put,
path = "/config/extension",
request_body = ExtensionQuery,
responses(
(status = 200, description = "Extension configuration updated successfully", body = String),
(status = 404, description = "Extension not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn update_extension(
State(state): State<AppState>,
headers: HeaderMap,
Json(extension): Json<ExtensionQuery>,
) -> Result<Json<String>, StatusCode> {
// Use the helper function to verify the secret key
verify_secret_key(&headers, &state)?;
let config = Config::global();
// Get current extensions
let mut extensions: HashMap<String, Value> = match config.get_param("extensions") {
Ok(exts) => exts,
Err(_) => return Err(StatusCode::NOT_FOUND),
};
// Check if extension exists
if !extensions.contains_key(&extension.name) {
return Err(StatusCode::NOT_FOUND);
}
// Update extension configuration
extensions.insert(extension.name.clone(), extension.config);
// Save updated extensions
match config.set_param(
"extensions",
Value::Object(extensions.into_iter().collect()),
) {
Ok(_) => Ok(Json(format!("Updated extension {}", extension.name))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
// Modified providers function using the new response type
#[utoipa::path(
get,
@@ -341,9 +380,11 @@ pub fn routes(state: AppState) -> Router {
.route("/config/upsert", post(upsert_config))
.route("/config/remove", post(remove_config))
.route("/config/read", post(read_config))
.route("/config/extension", post(add_extension))
.route("/config/extension", put(update_extension))
.route("/config/extension", delete(remove_extension))
.route("/config/extensions", get(get_extensions))
.route("/config/extensions", post(add_extension))
.route("/config/extensions/:name", put(update_extension))
.route("/config/extensions/:name", delete(remove_extension))
.route("/extensions/:name/toggle", post(toggle_extension))
.route("/config/providers", get(providers))
.with_state(state)
}
crates/goose/Cargo.toml
+1 -1
View File
@@ -60,7 +60,7 @@ serde_yaml = "0.9.34"
once_cell = "1.20.2"
etcetera = "0.8.0"
rand = "0.8.5"
utoipa = { version = "4.1" }
utoipa = "4.1"
# For Bedrock provider
aws-config = { version = "1.1.7", features = ["behavior-version-latest"] }
crates/goose/src/agents/capabilities.rs
+1 -1
View File
@@ -171,7 +171,7 @@ impl Capabilities {
.await
.map_err(|e| ExtensionError::Initialization(config.clone(), e))?;
let sanitized_name = normalize(config.name().to_string());
let sanitized_name = normalize(config.key().to_string());
// Store instructions if provided
if let Some(instructions) = init_result.instructions {
crates/goose/src/agents/extension.rs
+11 -3
View File
@@ -3,8 +3,10 @@ use std::collections::HashMap;
use mcp_client::client::Error as ClientError;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use utoipa::ToSchema;
use crate::config;
use crate::config::extensions::name_to_key;
/// Errors from Extension operation
#[derive(Error, Debug)]
@@ -21,7 +23,7 @@ pub enum ExtensionError {
pub type ExtensionResult<T> = Result<T, ExtensionError>;
#[derive(Debug, Clone, Deserialize, Serialize, Default)]
#[derive(Debug, Clone, Deserialize, Serialize, Default, ToSchema)]
pub struct Envs {
/// A map of environment variables to set, e.g. API_KEY -> some_secret, HOST -> host
#[serde(default)]
@@ -43,7 +45,7 @@ impl Envs {
}
/// Represents the different types of MCP extensions that can be added to the manager
#[derive(Debug, Clone, Deserialize, Serialize)]
#[derive(Debug, Clone, Deserialize, Serialize, ToSchema)]
#[serde(tag = "type")]
pub enum ExtensionConfig {
/// Server-sent events client with a URI endpoint
@@ -130,13 +132,19 @@ impl ExtensionConfig {
}
}
pub fn key(&self) -> String {
let name = self.name();
name_to_key(&name)
}
/// Get the extension name regardless of variant
pub fn name(&self) -> &str {
pub fn name(&self) -> String {
match self {
Self::Sse { name, .. } => name,
Self::Stdio { name, .. } => name,
Self::Builtin { name, .. } => name,
}
.to_string()
}
}
crates/goose/src/config/extensions.rs
+23 -15
View File
@@ -3,23 +3,28 @@ use crate::agents::ExtensionConfig;
use anyhow::Result;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use utoipa::ToSchema;
pub const DEFAULT_EXTENSION: &str = "developer";
pub const DEFAULT_EXTENSION_TIMEOUT: u64 = 300;
#[derive(Debug, Deserialize, Serialize, Clone)]
#[derive(Debug, Deserialize, Serialize, Clone, ToSchema)]
pub struct ExtensionEntry {
pub enabled: bool,
#[serde(flatten)]
pub config: ExtensionConfig,
}
pub fn name_to_key(name: &str) -> String {
name.to_string()
}
/// Extension configuration management
pub struct ExtensionManager;
impl ExtensionManager {
/// Get the extension configuration if enabled
pub fn get_config(name: &str) -> Result<Option<ExtensionConfig>> {
/// Get the extension configuration if enabled -- uses key
pub fn get_config(key: &str) -> Result<Option<ExtensionConfig>> {
let config = Config::global();
// Try to get the extension entry
@@ -28,7 +33,7 @@ impl ExtensionManager {
Err(super::ConfigError::NotFound(_)) => {
// Initialize with default developer extension
let defaults = HashMap::from([(
DEFAULT_EXTENSION.to_string(),
name_to_key(DEFAULT_EXTENSION), // Use key format for top-level key in config
ExtensionEntry {
enabled: true,
config: ExtensionConfig::Builtin {
@@ -43,7 +48,7 @@ impl ExtensionManager {
Err(e) => return Err(e.into()),
};
Ok(extensions.get(name).and_then(|entry| {
Ok(extensions.get(key).and_then(|entry| {
if entry.enabled {
Some(entry.config.clone())
} else {
@@ -60,33 +65,35 @@ impl ExtensionManager {
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
extensions.insert(entry.config.name().parse()?, entry);
let key = entry.config.key();
extensions.insert(key, entry);
config.set_param("extensions", serde_json::to_value(extensions)?)?;
Ok(())
}
/// Remove an extension configuration
pub fn remove(name: &str) -> Result<()> {
/// Remove an extension configuration -- uses the key
pub fn remove(key: &str) -> Result<()> {
let config = Config::global();
let mut extensions: HashMap<String, ExtensionEntry> = config
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
extensions.remove(name);
extensions.remove(key);
config.set_param("extensions", serde_json::to_value(extensions)?)?;
Ok(())
}
/// Enable or disable an extension
pub fn set_enabled(name: &str, enabled: bool) -> Result<()> {
/// Enable or disable an extension -- uses key
pub fn set_enabled(key: &str, enabled: bool) -> Result<()> {
let config = Config::global();
let mut extensions: HashMap<String, ExtensionEntry> = config
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
if let Some(entry) = extensions.get_mut(name) {
if let Some(entry) = extensions.get_mut(key) {
entry.enabled = enabled;
config.set_param("extensions", serde_json::to_value(extensions)?)?;
}
@@ -109,16 +116,17 @@ impl ExtensionManager {
.unwrap_or_else(|_| get_keys(Default::default())))
}
/// Check if an extension is enabled
pub fn is_enabled(name: &str) -> Result<bool> {
/// Check if an extension is enabled - FIXED to use key
pub fn is_enabled(key: &str) -> Result<bool> {
let config = Config::global();
let extensions: HashMap<String, ExtensionEntry> = config
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
Ok(extensions.get(name).map(|e| e.enabled).unwrap_or(false))
Ok(extensions.get(key).map(|e| e.enabled).unwrap_or(false))
}
}
fn get_keys(entries: HashMap<String, ExtensionEntry>) -> Vec<String> {
entries.into_keys().collect()
}
crates/goose/src/config/mod.rs
+1 -1
View File
@@ -1,6 +1,6 @@
mod base;
mod experiments;
mod extensions;
pub mod extensions;
pub use crate::agents::ExtensionConfig;
pub use base::{Config, ConfigError, APP_STRATEGY};