Просмотр исходного кода

feat(McpHub): inject env vars on whole mcp config (#3970)

* types(utils/config): improve return type infer

* feat(McpHub): inject env variable on whole mcp config

* fix: check for undefined configInjected

* refactor: improve type assertion

---------

Co-authored-by: Daniel <[email protected]>
Trung Dang 7 месяцев назад
Родитель
Сommit
dde71d24e4
2 измененных файлов с 16 добавлено и 13 удалено
  1. 13 10
      src/services/mcp/McpHub.ts
  2. 3 3
      src/utils/config.ts

+ 13 - 10
src/services/mcp/McpHub.ts

@@ -443,13 +443,16 @@ export class McpHub {
 
 			let transport: StdioClientTransport | SSEClientTransport
 
-			if (config.type === "stdio") {
+			// Inject environment variables to the config
+			const configInjected = (await injectEnv(config)) as typeof config
+
+			if (configInjected.type === "stdio") {
 				transport = new StdioClientTransport({
-					command: config.command,
-					args: config.args,
-					cwd: config.cwd,
+					command: configInjected.command,
+					args: configInjected.args,
+					cwd: configInjected.cwd,
 					env: {
-						...(config.env ? await injectEnv(config.env) : {}),
+					...(configInjected.env || {}),
 						...(process.env.PATH ? { PATH: process.env.PATH } : {}),
 						...(process.env.HOME ? { HOME: process.env.HOME } : {}),
 					},
@@ -508,16 +511,16 @@ export class McpHub {
 				// SSE connection
 				const sseOptions = {
 					requestInit: {
-						headers: config.headers,
+						headers: configInjected.headers,
 					},
 				}
 				// Configure ReconnectingEventSource options
 				const reconnectingEventSourceOptions = {
 					max_retry_time: 5000, // Maximum retry time in milliseconds
-					withCredentials: config.headers?.["Authorization"] ? true : false, // Enable credentials if Authorization header exists
+					withCredentials: configInjected.headers?.["Authorization"] ? true : false, // Enable credentials if Authorization header exists
 				}
 				global.EventSource = ReconnectingEventSource
-				transport = new SSEClientTransport(new URL(config.url), {
+				transport = new SSEClientTransport(new URL(configInjected.url), {
 					...sseOptions,
 					eventSourceInit: reconnectingEventSourceOptions,
 				})
@@ -537,9 +540,9 @@ export class McpHub {
 			const connection: McpConnection = {
 				server: {
 					name,
-					config: JSON.stringify(config),
+					config: JSON.stringify(configInjected),
 					status: "connecting",
-					disabled: config.disabled,
+					disabled: configInjected.disabled,
 					source,
 					projectPath: source === "project" ? vscode.workspace.workspaceFolders?.[0]?.uri.fsPath : undefined,
 					errorHistory: [],

+ 3 - 3
src/utils/config.ts

@@ -5,11 +5,11 @@
  *
  * Does not mutate original object
  */
-export async function injectEnv(config: string | Record<PropertyKey, any>, notFoundValue: any = "") {
+export async function injectEnv<C extends string | Record<PropertyKey, any>>(config: C, notFoundValue: any = "") {
 	// Use simple regex replace for now, will see if object traversal and recursion is needed here (e.g: for non-serializable objects)
 
 	const isObject = typeof config === "object"
-	let _config = isObject ? JSON.stringify(config) : config
+	let _config: string = isObject ? JSON.stringify(config) : config
 
 	_config = _config.replace(/\$\{env:([\w]+)\}/g, (_, name) => {
 		// Check if null or undefined
@@ -21,5 +21,5 @@ export async function injectEnv(config: string | Record<PropertyKey, any>, notFo
 		return process.env[name] ?? notFoundValue
 	})
 
-	return isObject ? JSON.parse(_config) : _config
+	return (isObject ? JSON.parse(_config) : _config) as C extends string ? string : C
 }