From 571ba86157725c177f087e8acc4f4ecbe5d8343e Mon Sep 17 00:00:00 2001
From: Alex Hancock <alexhancock@block.xyz>
Date: Thu, 17 Apr 2025 16:57:11 -0400
Subject: [PATCH] fix: specify extension env vars only by name (#2249)

Co-authored-by: Lily Delalande <ldelalande@block.xyz>
---
 .../settings_v2/extensions/agent-api.ts       | 43 +++++++++++++++++--
 .../extensions/extension-manager.ts           | 14 ++++++
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/ui/desktop/src/components/settings_v2/extensions/agent-api.ts b/ui/desktop/src/components/settings_v2/extensions/agent-api.ts
index 243a91c2..5dd9e74a 100644
--- a/ui/desktop/src/components/settings_v2/extensions/agent-api.ts
+++ b/ui/desktop/src/components/settings_v2/extensions/agent-api.ts
@@ -2,6 +2,8 @@ import { ExtensionConfig } from '../../../api/types.gen';
 import { getApiUrl, getSecretKey } from '../../../config';
 import { toastService, ToastServiceOptions } from '../../../toasts';
 import { replaceWithShims } from './utils';
+import { saveEnvVarsToKeyring } from './extension-manager';
+import type { AgentExtensionConfig } from './extension-manager';
 
 interface ApiResponse {
   error?: boolean;
@@ -141,13 +143,17 @@ export async function addToAgent(
   options: ToastServiceOptions = {}
 ): Promise<Response> {
   try {
-    if (extension.type === 'stdio') {
-      extension.cmd = await replaceWithShims(extension.cmd);
+    await saveEnvVarsToKeyring(extension);
+
+    const ext = toAgentExtensionConfig(extension);
+
+    if (ext.type === 'stdio') {
+      ext.cmd = await replaceWithShims(ext.cmd);
     }
 
-    extension.name = sanitizeName(extension.name);
+    ext.name = sanitizeName(ext.name);
 
-    return await extensionApiCall('/extensions/add', extension, options);
+    return await extensionApiCall('/extensions/add', ext, options);
   } catch (error) {
     // Check if this is a 428 error and make the message more descriptive
     if (error.message && error.message.includes('428')) {
@@ -179,3 +185,32 @@ export async function removeFromAgent(
 function sanitizeName(name: string) {
   return name.toLowerCase().replace(/-/g, '').replace(/_/g, '').replace(/\s/g, '');
 }
+
+export function toAgentExtensionConfig(config: ExtensionConfig): AgentExtensionConfig {
+  // Use type narrowing to handle different variants of the union type
+  if ('type' in config) {
+    switch (config.type) {
+      case 'sse': {
+        const { envs, ...rest } = config;
+        return {
+          ...rest,
+          env_keys: envs ? Object.keys(envs) : undefined,
+        };
+      }
+      case 'stdio': {
+        const { envs, ...rest } = config;
+        return {
+          ...rest,
+          env_keys: envs ? Object.keys(envs) : undefined,
+        };
+      }
+      case 'builtin':
+      case 'frontend':
+        // These types don't have envs field, so just return as is
+        return config;
+    }
+  }
+
+  // This should never happen due to the union type constraint
+  throw new Error('Invalid extension configuration type');
+}
diff --git a/ui/desktop/src/components/settings_v2/extensions/extension-manager.ts b/ui/desktop/src/components/settings_v2/extensions/extension-manager.ts
index bd0f4113..c7d8496e 100644
--- a/ui/desktop/src/components/settings_v2/extensions/extension-manager.ts
+++ b/ui/desktop/src/components/settings_v2/extensions/extension-manager.ts
@@ -1,6 +1,12 @@
 import type { ExtensionConfig } from '../../../api/types.gen';
 import { toastService, ToastServiceOptions } from '../../../toasts';
 import { addToAgent, removeFromAgent } from './agent-api';
+import { upsertConfig } from '../../../api';
+
+// TODO: unify config.yaml and the agent /extensions/add API's notion of env vars
+export type AgentExtensionConfig = ExtensionConfig & {
+  env_keys?: string[];
+};
 
 interface ActivateExtensionProps {
   addToConfig: (name: string, extensionConfig: ExtensionConfig, enabled: boolean) => Promise<void>;
@@ -283,3 +289,11 @@ export async function deleteExtension({ name, removeFromConfig }: DeleteExtensio
     throw agentRemoveError;
   }
 }
+
+export async function saveEnvVarsToKeyring(extension: ExtensionConfig) {
+  if (extension.type === 'stdio' || extension.type === 'sse') {
+    for (const [key, value] of Object.entries(extension.envs || {})) {
+      await upsertConfig({ body: { key, value, is_secret: true } });
+    }
+  }
+}