tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 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-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)
}