fix: use env keys (#2258)

Co-authored-by: Zaki Ali <zaki@squareup.com>
Co-authored-by: Kalvin C <kalvinnchau@users.noreply.github.com>

crates/goose/src/agents/extension.rs

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