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

feat: endpoints for config management via goose-server (#1207)

Browse Source 
Co-authored-by: Lily Delalande <ldelalande@squareup.com>
This commit is contained in:
Alex Hancock
2025-02-12 21:19:13 -05:00
committed by GitHub
parent e461e6845d
commit 9287eaeb5d
14 changed files with 569 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/bin/generate_schema.rs
+22
View File
@@ -0,0 +1,22 @@
use goose_server::openapi;
use std::env;
use std::fs;
fn main() {
let schema = openapi::generate_schema();
// Get the current working directory
let current_dir = env::current_dir().unwrap();
let output_path = current_dir.join("ui").join("desktop").join("openapi.json");
// Ensure parent directory exists
if let Some(parent) = output_path.parent() {
fs::create_dir_all(parent).unwrap();
}
fs::write(&output_path, schema).unwrap();
println!(
"Successfully generated OpenAPI schema at {}",
output_path.display()
);
}
crates/goose-server/src/lib.rs
+7
View File
@@ -0,0 +1,7 @@
pub mod openapi;
pub mod routes;
pub mod state;
// Re-export commonly used items
pub use openapi::*;
pub use state::*;
crates/goose-server/src/main.rs
+1
View File
@@ -11,6 +11,7 @@ mod commands;
mod configuration;
mod error;
mod logging;
mod openapi;
mod routes;
mod state;
crates/goose-server/src/openapi.rs
+27
View File
@@ -0,0 +1,27 @@
use utoipa::OpenApi;
#[allow(dead_code)] // Used by utoipa for OpenAPI generation
#[derive(OpenApi)]
#[openapi(
paths(
super::routes::config_management::upsert_config,
super::routes::config_management::remove_config,
super::routes::config_management::read_config,
super::routes::config_management::add_extension,
super::routes::config_management::remove_extension,
super::routes::config_management::read_all_config
),
components(schemas(
super::routes::config_management::UpsertConfigQuery,
super::routes::config_management::ConfigKeyQuery,
super::routes::config_management::ExtensionQuery,
super::routes::config_management::ConfigResponse
))
)]
pub struct ApiDoc;
#[allow(dead_code)] // Used by generate_schema binary
pub fn generate_schema() -> String {
let api_doc = ApiDoc::openapi();
serde_json::to_string_pretty(&api_doc).unwrap()
}
crates/goose-server/src/routes/agent.rs
+1
View File
@@ -95,6 +95,7 @@ async fn extend_prompt(
}
}
#[axum::debug_handler]
async fn create_agent(
State(state): State<AppState>,
headers: HeaderMap,
crates/goose-server/src/routes/config_management.rs
+206
View File
@@ -0,0 +1,206 @@
use axum::{
extract::State,
routing::{delete, get, post},
Json, Router,
};
use goose::config::Config;
use http::StatusCode;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::{collections::HashMap, sync::Arc};
use tokio::sync::Mutex;
use utoipa::ToSchema;
use crate::state::AppState;
#[derive(Deserialize, ToSchema)]
pub struct UpsertConfigQuery {
pub key: String,
pub value: Value,
pub is_secret: Option<bool>,
}
#[derive(Deserialize, ToSchema)]
pub struct ConfigKeyQuery {
pub key: String,
}
#[derive(Deserialize, ToSchema)]
pub struct ExtensionQuery {
pub name: String,
pub config: Value,
}
#[derive(Serialize, ToSchema)]
pub struct ConfigResponse {
pub config: HashMap<String, Value>,
}
#[utoipa::path(
post,
path = "/config/upsert",
request_body = UpsertConfigQuery,
responses(
(status = 200, description = "Configuration value upserted successfully", body = String),
(status = 500, description = "Internal server error")
)
)]
pub async fn upsert_config(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
Json(query): Json<UpsertConfigQuery>,
) -> Result<Json<Value>, StatusCode> {
let config = Config::global();
let result = if query.is_secret.unwrap_or(false) {
config.set_secret(&query.key, query.value)
} else {
config.set(&query.key, query.value)
};
match result {
Ok(_) => Ok(Json(Value::String(format!("Upserted key {}", query.key)))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
#[utoipa::path(
post,
path = "/config/remove",
request_body = ConfigKeyQuery,
responses(
(status = 200, description = "Configuration value removed successfully", body = String),
(status = 404, description = "Configuration key not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn remove_config(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
Json(query): Json<ConfigKeyQuery>,
) -> Result<Json<String>, StatusCode> {
let config = Config::global();
match config.delete(&query.key) {
Ok(_) => Ok(Json(format!("Removed key {}", query.key))),
Err(_) => Err(StatusCode::NOT_FOUND),
}
}
#[utoipa::path(
get,
path = "/config/read",
request_body = ConfigKeyQuery,
responses(
(status = 200, description = "Configuration value retrieved successfully", body = Value),
(status = 404, description = "Configuration key not found")
)
)]
pub async fn read_config(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
Json(query): Json<ConfigKeyQuery>,
) -> Result<Json<Value>, StatusCode> {
let config = Config::global();
match config.get::<Value>(&query.key) {
Ok(value) => Ok(Json(value)),
Err(_) => Err(StatusCode::NOT_FOUND),
}
}
#[utoipa::path(
post,
path = "/config/extension",
request_body = ExtensionQuery,
responses(
(status = 200, description = "Extension added successfully", body = String),
(status = 400, description = "Invalid request"),
(status = 500, description = "Internal server error")
)
)]
pub async fn add_extension(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
Json(extension): Json<ExtensionQuery>,
) -> Result<Json<String>, StatusCode> {
let config = Config::global();
// Get current extensions or initialize empty map
let mut extensions: HashMap<String, Value> =
config.get("extensions").unwrap_or_else(|_| HashMap::new());
// Add new extension
extensions.insert(extension.name.clone(), extension.config);
// Save updated extensions
match config.set(
"extensions",
Value::Object(extensions.into_iter().collect()),
) {
Ok(_) => Ok(Json(format!("Added extension {}", extension.name))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
}
#[utoipa::path(
delete,
path = "/config/extension",
request_body = ConfigKeyQuery,
responses(
(status = 200, description = "Extension removed successfully", body = String),
(status = 404, description = "Extension not found"),
(status = 500, description = "Internal server error")
)
)]
pub async fn remove_extension(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
Json(query): Json<ConfigKeyQuery>,
) -> Result<Json<String>, StatusCode> {
let config = Config::global();
// Get current extensions
let mut extensions: HashMap<String, Value> = match config.get("extensions") {
Ok(exts) => exts,
Err(_) => return Err(StatusCode::NOT_FOUND),
};
// Remove extension if it exists
if extensions.remove(&query.key).is_some() {
// Save updated extensions
match config.set(
"extensions",
Value::Object(extensions.into_iter().collect()),
) {
Ok(_) => Ok(Json(format!("Removed extension {}", query.key))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
} else {
Err(StatusCode::NOT_FOUND)
}
}
#[utoipa::path(
get,
path = "/config",
responses(
(status = 200, description = "All configuration values retrieved successfully", body = ConfigResponse)
)
)]
pub async fn read_all_config(
State(_state): State<Arc<Mutex<HashMap<String, Value>>>>,
) -> Result<Json<ConfigResponse>, StatusCode> {
let config = Config::global();
// Load values from config file
let values = config.load_values().unwrap_or_default();
Ok(Json(ConfigResponse { config: values }))
}
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/config", get(read_all_config))
.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", delete(remove_extension))
.with_state(state.config)
}
crates/goose-server/src/routes/mod.rs
+3 -1
View File
@@ -1,5 +1,6 @@
// Export route modules
pub mod agent;
pub mod config_management;
pub mod configs;
pub mod extension;
pub mod health;
@@ -14,5 +15,6 @@ pub fn configure(state: crate::state::AppState) -> Router {
.merge(reply::routes(state.clone()))
.merge(agent::routes(state.clone()))
.merge(extension::routes(state.clone()))
.merge(configs::routes(state))
.merge(configs::routes(state.clone()))
.merge(config_management::routes(state))
}
crates/goose-server/src/routes/reply.rs
+2
View File
@@ -559,6 +559,7 @@ mod tests {
mod integration_tests {
use super::*;
use axum::{body::Body, http::Request};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
use tower::ServiceExt;
@@ -573,6 +574,7 @@ mod tests {
});
let agent = AgentFactory::create("reference", mock_provider).unwrap();
let state = AppState {
config: Arc::new(Mutex::new(HashMap::new())), // Add this line
agent: Arc::new(Mutex::new(Some(agent))),
secret_key: "test-secret".to_string(),
};
crates/goose-server/src/state.rs
+4
View File
@@ -1,5 +1,7 @@
use anyhow::Result;
use goose::agents::Agent;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::Mutex;
@@ -9,6 +11,7 @@ use tokio::sync::Mutex;
pub struct AppState {
pub agent: Arc<Mutex<Option<Box<dyn Agent>>>>,
pub secret_key: String,
pub config: Arc<Mutex<HashMap<String, Value>>>,
}
impl AppState {
@@ -16,6 +19,7 @@ impl AppState {
Ok(Self {
agent: Arc::new(Mutex::new(None)),
secret_key,
config: Arc::new(Mutex::new(HashMap::new())),
})
}
}