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

fix: use env keys (#2258)

Browse Source 
Co-authored-by: Zaki Ali <zaki@squareup.com>
Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com>
This commit is contained in:
Bradley Axen
2025-04-18 14:01:46 -07:00
committed by GitHub
parent 621eb42fb0
commit cfb0eab9cf
19 changed files with 506 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/openapi.rs
+1
View File
@@ -12,6 +12,7 @@ use utoipa::OpenApi;
#[derive(OpenApi)]
#[openapi(
paths(
super::routes::config_management::backup_config,
super::routes::config_management::init_config,
super::routes::config_management::upsert_config,
super::routes::config_management::remove_config,
crates/goose-server/src/routes/config_management.rs
+52 -1
View File
@@ -292,7 +292,9 @@ pub async fn read_all_config(
let config = Config::global();
// Load values from config file
let values = config.load_values().unwrap_or_default();
let values = config
.load_values()
.map_err(|_| StatusCode::UNPROCESSABLE_ENTITY)?;
Ok(Json(ConfigResponse { config: values }))
}
@@ -429,6 +431,54 @@ pub async fn upsert_permissions(
Ok(Json("Permissions updated successfully".to_string()))
}
use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs};
use once_cell::sync::Lazy;
pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
top_level_domain: "Block".to_string(),
author: "Block".to_string(),
app_name: "goose".to_string(),
});
#[utoipa::path(
post,
path = "/config/backup",
responses(
(status = 200, description = "Config file backed up", body = String),
(status = 500, description = "Internal server error")
)
)]
pub async fn backup_config(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<String>, StatusCode> {
verify_secret_key(&headers, &state)?;
let config_dir = choose_app_strategy(APP_STRATEGY.clone())
.expect("goose requires a home dir")
.config_dir();
let config_path = config_dir.join("config.yaml");
if config_path.exists() {
let file_name = config_path
.file_name()
.ok_or(StatusCode::INTERNAL_SERVER_ERROR)?;
// Append ".bak" to the file name
let mut backup_name = file_name.to_os_string();
backup_name.push(".bak");
// Construct the new path with the same parent directory
let backup = config_path.with_file_name(backup_name);
match std::fs::rename(&config_path, &backup) {
Ok(_) => Ok(Json(format!("Moved {:?} to {:?}", config_path, backup))),
Err(_) => Err(StatusCode::INTERNAL_SERVER_ERROR),
}
} else {
Err(StatusCode::INTERNAL_SERVER_ERROR)
}
}
pub fn routes(state: AppState) -> Router {
Router::new()
.route("/config", get(read_all_config))
@@ -440,6 +490,7 @@ pub fn routes(state: AppState) -> Router {
.route("/config/extensions/:name", delete(remove_extension))
.route("/config/providers", get(providers))
.route("/config/init", post(init_config))
.route("/config/backup", post(backup_config))
.route("/config/permissions", post(upsert_permissions))
.with_state(state)
}
crates/goose-server/src/routes/extension.rs
+20 -66
View File
@@ -1,14 +1,10 @@
use std::collections::HashMap;
use std::env;
use std::path::Path;
use std::sync::OnceLock;
use crate::state::AppState;
use axum::{extract::State, routing::post, Json, Router};
use goose::{
agents::{extension::Envs, ExtensionConfig},
config::Config,
};
use goose::agents::{extension::Envs, ExtensionConfig};
use http::{HeaderMap, StatusCode};
use serde::{Deserialize, Serialize};
use tracing;
@@ -24,6 +20,9 @@ enum ExtensionConfigRequest {
name: String,
/// The URI endpoint for the SSE extension.
uri: String,
#[serde(default)]
/// Map of environment variable key to values.
envs: Envs,
/// List of environment variable keys. The server will fetch their values from the keyring.
#[serde(default)]
env_keys: Vec<String>,
@@ -39,6 +38,9 @@ enum ExtensionConfigRequest {
/// Arguments for the command.
#[serde(default)]
args: Vec<String>,
#[serde(default)]
/// Map of environment variable key to values.
envs: Envs,
/// List of environment variable keys. The server will fetch their values from the keyring.
#[serde(default)]
env_keys: Vec<String>,
@@ -162,55 +164,28 @@ async fn add_extension(
}
}
// Load the configuration
let config = Config::global();
// Initialize a vector to collect any missing keys.
let mut missing_keys = Vec::new();
// Construct ExtensionConfig with Envs populated from keyring based on provided env_keys.
let extension_config: ExtensionConfig = match request {
ExtensionConfigRequest::Sse {
name,
uri,
envs,
env_keys,
timeout,
} => {
let mut env_map = HashMap::new();
for key in env_keys {
match config.get_secret(&key) {
Ok(value) => {
env_map.insert(key, value);
}
Err(_) => {
missing_keys.push(key);
}
}
}
if !missing_keys.is_empty() {
return Ok(Json(ExtensionResponse {
error: true,
message: Some(format!(
"Missing secrets for keys: {}",
missing_keys.join(", ")
)),
}));
}
ExtensionConfig::Sse {
name,
uri,
envs: Envs::new(env_map),
description: None,
timeout,
bundled: None,
}
}
} => ExtensionConfig::Sse {
name,
uri,
envs,
env_keys,
description: None,
timeout,
bundled: None,
},
ExtensionConfigRequest::Stdio {
name,
cmd,
args,
envs,
env_keys,
timeout,
} => {
@@ -226,34 +201,13 @@ async fn add_extension(
// }));
// }
let mut env_map = HashMap::new();
for key in env_keys {
match config.get_secret(&key) {
Ok(value) => {
env_map.insert(key, value);
}
Err(_) => {
missing_keys.push(key);
}
}
}
if !missing_keys.is_empty() {
return Ok(Json(ExtensionResponse {
error: true,
message: Some(format!(
"Missing secrets for keys: {}",
missing_keys.join(", ")
)),
}));
}
ExtensionConfig::Stdio {
name,
cmd,
args,
description: None,
envs: Envs::new(env_map),
envs,
env_keys,
timeout,
bundled: None,
}