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

Co-authored-by: Lily Delalande <ldelalande@squareup.com>

crates/goose-server/Cargo.toml

@@ -12,7 +12,7 @@ goose = { path = "../goose" }
 mcp-core = { path = "../mcp-core" }
 goose-mcp = { path = "../goose-mcp" }
 mcp-server = { path = "../mcp-server" }
-axum = { version = "0.7", features = ["ws"] }
+axum = { version = "0.7.2", features = ["ws", "macros"] }
 tokio = { version = "1.0", features = ["full"] }
 chrono = "0.4"
 tower-http = { version = "0.5", features = ["cors"] }
@@ -31,11 +31,19 @@ thiserror = "1.0"
 clap = { version = "4.4", features = ["derive"] }
 once_cell = "1.20.2"
 etcetera = "0.8.0"
+serde_yaml = "0.9.34"
+axum-extra = "0.10.0"
+utoipa = { version = "4.1", features = ["axum_extras"] }
+dirs = "6.0.0"
 
 [[bin]]
 name = "goosed"
 path = "src/main.rs"
 
+[[bin]]
+name = "generate_schema"
+path = "src/bin/generate_schema.rs"
+
 [dev-dependencies]
 tower = "0.5"
-async-trait = "0.1"
+async-trait = "0.1"

crates/goose-server/build.rs

@@ -0,0 +1,4 @@
+// We'll generate the schema at runtime since we need access to the complete application context
+fn main() {
+    println!("cargo:rerun-if-changed=src/");
+}

crates/goose-server/src/bin/generate_schema.rs

@@ -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

@@ -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

@@ -11,6 +11,7 @@ mod commands;
 mod configuration;
 mod error;
 mod logging;
+mod openapi;
 mod routes;
 mod state;
 

crates/goose-server/src/openapi.rs

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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,13 +11,15 @@ 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 {
     pub async fn new(secret_key: String) -> Result<Self> {
         Ok(Self {
             agent: Arc::new(Mutex::new(None)),
             secret_key,
+            config: Arc::new(Mutex::new(HashMap::new())),
         })
     }
 }