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-cli/src/commands/configure.rs
+32 -2
View File
@@ -583,6 +583,9 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
cliclack::confirm("Would you like to add environment variables?").interact()?;
let mut envs = HashMap::new();
let mut env_keys = Vec::new();
let config = Config::global();
if add_env {
loop {
let key: String = cliclack::input("Environment variable name:")
@@ -593,7 +596,18 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
.mask('▪')
.interact()?;
envs.insert(key, value);
// Try to store in keychain
let keychain_key = key.to_string();
match config.set_secret(&keychain_key, Value::String(value.clone())) {
Ok(_) => {
// Successfully stored in keychain, add to env_keys
env_keys.push(keychain_key);
}
Err(_) => {
// Failed to store in keychain, store directly in envs
envs.insert(key, value);
}
}
if !cliclack::confirm("Add another environment variable?").interact()? {
break;
@@ -608,6 +622,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
cmd,
args,
envs: Envs::new(envs),
env_keys,
description,
timeout: Some(timeout),
bundled: None,
@@ -671,6 +686,9 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
cliclack::confirm("Would you like to add environment variables?").interact()?;
let mut envs = HashMap::new();
let mut env_keys = Vec::new();
let config = Config::global();
if add_env {
loop {
let key: String = cliclack::input("Environment variable name:")
@@ -681,7 +699,18 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
.mask('▪')
.interact()?;
envs.insert(key, value);
// Try to store in keychain
let keychain_key = key.to_string();
match config.set_secret(&keychain_key, Value::String(value.clone())) {
Ok(_) => {
// Successfully stored in keychain, add to env_keys
env_keys.push(keychain_key);
}
Err(_) => {
// Failed to store in keychain, store directly in envs
envs.insert(key, value);
}
}
if !cliclack::confirm("Add another environment variable?").interact()? {
break;
@@ -695,6 +724,7 @@ pub fn configure_extensions_dialog() -> Result<(), Box<dyn Error>> {
name: name.clone(),
uri,
envs: Envs::new(envs),
env_keys,
description,
timeout: Some(timeout),
bundled: None,
crates/goose-cli/src/session/mod.rs
+2
View File
@@ -158,6 +158,7 @@ impl Session {
cmd,
args: parts.iter().map(|s| s.to_string()).collect(),
envs: Envs::new(envs),
env_keys: Vec::new(),
description: Some(goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string()),
// TODO: should set timeout
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
@@ -190,6 +191,7 @@ impl Session {
name,
uri: extension_url,
envs: Envs::new(HashMap::new()),
env_keys: Vec::new(),
description: Some(goose::config::DEFAULT_EXTENSION_DESCRIPTION.to_string()),
// TODO: should set timeout
timeout: Some(goose::config::DEFAULT_EXTENSION_TIMEOUT),
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,
}
crates/goose/src/agents/extension.rs
+10
View File
@@ -24,6 +24,8 @@ pub enum ExtensionError {
Transport(#[from] mcp_client::transport::Error),
#[error("Environment variable `{0}` is not allowed to be overridden.")]
InvalidEnvVar(String),
#[error("Error during extension setup: {0}")]
SetupError(String),
#[error("Join error occurred during task execution: {0}")]
TaskJoinError(#[from] tokio::task::JoinError),
}
@@ -128,6 +130,8 @@ pub enum ExtensionConfig {
uri: String,
#[serde(default)]
envs: Envs,
#[serde(default)]
env_keys: Vec<String>,
description: Option<String>,
// NOTE: set timeout to be optional for compatibility.
// However, new configurations should include this field.
@@ -145,6 +149,8 @@ pub enum ExtensionConfig {
args: Vec<String>,
#[serde(default)]
envs: Envs,
#[serde(default)]
env_keys: Vec<String>,
timeout: Option<u64>,
description: Option<String>,
/// Whether this extension is bundled with Goose
@@ -194,6 +200,7 @@ impl ExtensionConfig {
name: name.into(),
uri: uri.into(),
envs: Envs::default(),
env_keys: Vec::new(),
description: Some(description.into()),
timeout: Some(timeout.into()),
bundled: None,
@@ -211,6 +218,7 @@ impl ExtensionConfig {
cmd: cmd.into(),
args: vec![],
envs: Envs::default(),
env_keys: Vec::new(),
description: Some(description.into()),
timeout: Some(timeout.into()),
bundled: None,
@@ -227,6 +235,7 @@ impl ExtensionConfig {
name,
cmd,
envs,
env_keys,
timeout,
description,
bundled,
@@ -235,6 +244,7 @@ impl ExtensionConfig {
name,
cmd,
envs,
env_keys,
args: args.into_iter().map(Into::into).collect(),
description,
timeout,
crates/goose/src/agents/extension_manager.rs
+71 -9
View File
@@ -10,10 +10,11 @@ use std::sync::LazyLock;
use std::time::Duration;
use tokio::sync::Mutex;
use tokio::task;
use tracing::debug;
use tracing::{debug, error, warn};
use super::extension::{ExtensionConfig, ExtensionError, ExtensionInfo, ExtensionResult, ToolInfo};
use crate::config::ExtensionConfigManager;
use crate::agents::extension::Envs;
use crate::config::{Config, ExtensionConfigManager};
use crate::prompt_template;
use mcp_client::client::{ClientCapabilities, ClientInfo, McpClient, McpClientTrait};
use mcp_client::transport::{SseTransport, StdioTransport, Transport};
@@ -113,11 +114,74 @@ impl ExtensionManager {
// TODO IMPORTANT need to ensure this times out if the extension command is broken!
pub async fn add_extension(&mut self, config: ExtensionConfig) -> ExtensionResult<()> {
let sanitized_name = normalize(config.key().to_string());
/// Helper function to merge environment variables from direct envs and keychain-stored env_keys
async fn merge_environments(
envs: &Envs,
env_keys: &[String],
ext_name: &str,
) -> Result<HashMap<String, String>, ExtensionError> {
let mut all_envs = envs.get_env();
let config_instance = Config::global();
for key in env_keys {
// If the Envs payload already contains the key, prefer that value
// over looking into the keychain/secret store
if all_envs.contains_key(key) {
continue;
}
match config_instance.get(key, true) {
Ok(value) => {
if value.is_null() {
warn!(
key = %key,
ext_name = %ext_name,
"Secret key not found in config (returned null)."
);
continue;
}
// Try to get string value
if let Some(str_val) = value.as_str() {
all_envs.insert(key.clone(), str_val.to_string());
} else {
warn!(
key = %key,
ext_name = %ext_name,
value_type = %value.get("type").and_then(|t| t.as_str()).unwrap_or("unknown"),
"Secret value is not a string; skipping."
);
}
}
Err(e) => {
error!(
key = %key,
ext_name = %ext_name,
error = %e,
"Failed to fetch secret from config."
);
return Err(ExtensionError::SetupError(format!(
"Failed to fetch secret '{}' from config: {}",
key, e
)));
}
}
}
Ok(all_envs)
}
let mut client: Box<dyn McpClientTrait> = match &config {
ExtensionConfig::Sse {
uri, envs, timeout, ..
uri,
envs,
env_keys,
timeout,
..
} => {
let transport = SseTransport::new(uri, envs.get_env());
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = SseTransport::new(uri, all_envs);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
@@ -131,10 +195,12 @@ impl ExtensionManager {
cmd,
args,
envs,
env_keys,
timeout,
..
} => {
let transport = StdioTransport::new(cmd, args.to_vec(), envs.get_env());
let all_envs = merge_environments(envs, env_keys, &sanitized_name).await?;
let transport = StdioTransport::new(cmd, args.to_vec(), all_envs);
let handle = transport.start().await?;
let service = McpService::with_timeout(
handle,
@@ -150,7 +216,6 @@ impl ExtensionManager {
timeout,
bundled: _,
} => {
// For builtin extensions, we run the current executable with mcp and extension name
let cmd = std::env::current_exe()
.expect("should find the current executable")
.to_str()
@@ -185,19 +250,16 @@ impl ExtensionManager {
.await
.map_err(|e| ExtensionError::Initialization(config.clone(), e))?;
// Store instructions if provided
if let Some(instructions) = init_result.instructions {
self.instructions
.insert(sanitized_name.clone(), instructions);
}
// if the server is capable if resources we track it
if init_result.capabilities.resources.is_some() {
self.resource_capable_extensions
.insert(sanitized_name.clone());
}
// Store the client using the provided name
self.clients
.insert(sanitized_name.clone(), Arc::new(Mutex::new(client)));
crates/goose/src/config/extensions.rs
+1 -3
View File
@@ -126,9 +126,7 @@ impl ExtensionConfigManager {
/// Get all extensions and their configurations
pub fn get_all() -> Result<Vec<ExtensionEntry>> {
let config = Config::global();
let extensions: HashMap<String, ExtensionEntry> = config
.get_param("extensions")
.unwrap_or_else(|_| HashMap::new());
let extensions: HashMap<String, ExtensionEntry> = config.get_param("extensions")?;
Ok(Vec::from_iter(extensions.values().cloned()))
}