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

feat: add `injectVariables()`, support magic variables for MCPs (`workspaceFolder`) (#4442)

* feat: add `injectVariables()`, support magic variables for MCPs

* fix: fallback for `workspaceFolder` should just be an empty string

Previously this is intended so that the CLI receives a correct empty path argument, but on a second thought, if the user have added the quotes themselves it might cause error.

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>

* chore: remove unused import

* chore: better log format

* chore: better describe the accepted config type and more extensive test

---------

Co-authored-by: ellipsis-dev[bot] <65095814+ellipsis-dev[bot]@users.noreply.github.com>
Trung Dang 6 месяцев назад
Родитель
Сommit
84ccf3fcc3
3 измененных файлов с 105 добавлено и 18 удалено
  1. 6 3
      src/services/mcp/McpHub.ts
  2. 55 5
      src/utils/__tests__/config.spec.ts
  3. 44 10
      src/utils/config.ts

+ 6 - 3
src/services/mcp/McpHub.ts

@@ -31,7 +31,7 @@ import {
 } from "../../shared/mcp"
 import { fileExistsAtPath } from "../../utils/fs"
 import { arePathsEqual } from "../../utils/path"
-import { injectEnv } from "../../utils/config"
+import { injectVariables } from "../../utils/config"
 
 export type McpConnection = {
 	server: McpServer
@@ -579,8 +579,11 @@ export class McpHub {
 
 			let transport: StdioClientTransport | SSEClientTransport | StreamableHTTPClientTransport
 
-			// Inject environment variables to the config
-			const configInjected = (await injectEnv(config)) as typeof config
+			// Inject variables to the config (environment, magic variables,...)
+			const configInjected = (await injectVariables(config, {
+				env: process.env,
+				workspaceFolder: vscode.workspace.workspaceFolders?.[0]?.uri.fsPath ?? "",
+			})) as typeof config
 
 			if (configInjected.type === "stdio") {
 				transport = new StdioClientTransport({

+ 55 - 5
src/utils/__tests__/config.spec.ts

@@ -1,5 +1,6 @@
 import { vitest, describe, it, expect, beforeEach, afterAll } from "vitest"
-import { injectEnv } from "../config"
+import { injectEnv, injectVariables } from "../config"
+
 
 describe("injectEnv", () => {
 	const originalEnv = process.env
@@ -30,14 +31,44 @@ describe("injectEnv", () => {
 			key: "${env:API_KEY}",
 			url: "${env:ENDPOINT}",
 			nested: {
-				value: "Keep this ${env:API_KEY}",
+				string: "Keep this ${env:API_KEY}",
+				number: 123,
+				boolean: true,
+				stringArr: ["${env:API_KEY}", "${env:ENDPOINT}"],
+				numberArr: [123, 456],
+				booleanArr: [true, false],
+			},
+			deeply: {
+				nested: {
+					string: "Keep this ${env:API_KEY}",
+					number: 123,
+					boolean: true,
+					stringArr: ["${env:API_KEY}", "${env:ENDPOINT}"],
+					numberArr: [123, 456],
+					booleanArr: [true, false],
+				},
 			},
 		}
 		const expectedObject = {
 			key: "12345",
 			url: "https://example.com",
 			nested: {
-				value: "Keep this 12345",
+				string: "Keep this 12345",
+				number: 123,
+				boolean: true,
+				stringArr: ["12345", "https://example.com"],
+				numberArr: [123, 456],
+				booleanArr: [true, false],
+			},
+			deeply: {
+				nested: {
+					string: "Keep this 12345",
+					number: 123,
+					boolean: true,
+					stringArr: ["12345", "https://example.com"],
+					numberArr: [123, 456],
+					booleanArr: [true, false],
+				},
 			},
 		}
 		const result = await injectEnv(configObject)
@@ -52,7 +83,7 @@ describe("injectEnv", () => {
 		const result = await injectEnv(configString, "NOT_FOUND")
 		expect(result).toBe(expectedString)
 		expect(consoleWarnSpy).toHaveBeenCalledWith(
-			"[injectEnv] env variable MISSING_VAR referenced but not found in process.env",
+			`[injectVariables] variable "MISSING_VAR" referenced but not found in "env"`,
 		)
 		consoleWarnSpy.mockRestore()
 	})
@@ -64,7 +95,7 @@ describe("injectEnv", () => {
 		const result = await injectEnv(configString)
 		expect(result).toBe(expectedString)
 		expect(consoleWarnSpy).toHaveBeenCalledWith(
-			"[injectEnv] env variable ANOTHER_MISSING referenced but not found in process.env",
+			`[injectVariables] variable "ANOTHER_MISSING" referenced but not found in "env"`,
 		)
 		consoleWarnSpy.mockRestore()
 	})
@@ -99,3 +130,22 @@ describe("injectEnv", () => {
 		expect(result).toEqual({})
 	})
 })
+
+describe("injectVariables", () => {
+	it("should replace singular variable", async () => {
+		const result = await injectVariables("Hello ${v}", { v: "Hola" })
+		expect(result).toEqual("Hello Hola")
+	})
+
+	it("should handle undefined singular variable input", async () => {
+		const result = await injectVariables("Hello ${v}", { v: undefined })
+		expect(result).toEqual("Hello ${v}")
+	})
+
+	it("should handle empty string singular variable input", async () => {
+		const result = await injectVariables("Hello ${v}", { v: "" })
+		expect(result).toEqual("Hello ")
+	})
+
+	// Variable maps are already tested by `injectEnv` tests above.
+})

+ 44 - 10
src/utils/config.ts

@@ -1,3 +1,15 @@
+export type InjectableConfigType =
+	| string
+	| {
+			[key: string]:
+				| undefined
+				| null
+				| boolean
+				| number
+				| InjectableConfigType
+				| Array<undefined | null | boolean | number | InjectableConfigType>
+	  }
+
 /**
  * Deeply injects environment variables into a configuration object/string/json
  *
@@ -5,21 +17,43 @@
  *
  * Does not mutate original object
  */
-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)
+export async function injectEnv<C extends InjectableConfigType>(config: C, notFoundValue: any = "") {
+	return injectVariables(config, { env: process.env }, notFoundValue)
+}
 
+/**
+ * Deeply injects variables into a configuration object/string/json
+ *
+ * Uses VSCode's variables reference pattern: https://code.visualstudio.com/docs/reference/variables-reference#_environment-variables
+ *
+ * Does not mutate original object
+ *
+ * There is a special handling for a nested (record-type) variables, where it is replaced by `propNotFoundValue` (if available) if the root key exists but the nested key does not.
+ *
+ * Matched keys that have `null` | `undefined` values are treated as not found.
+ */
+export async function injectVariables<C extends InjectableConfigType>(
+	config: C,
+	variables: Record<string, undefined | null | string | Record<string, undefined | null | string>>,
+	propNotFoundValue?: 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: string = isObject ? JSON.stringify(config) : config
 
-	_config = _config.replace(/\$\{env:([\w]+)\}/g, (_, name) => {
-		// Check if null or undefined
-		// intentionally using == to match null | undefined
-		if (process.env[name] == null) {
-			console.warn(`[injectEnv] env variable ${name} referenced but not found in process.env`)
-		}
+	// Intentionally using `== null` to match null | undefined
+	for (const [key, value] of Object.entries(variables)) {
+		if (value == null) continue
+
+		if (typeof value === "string") _config = _config.replace(new RegExp(`\\$\\{${key}\\}`, "g"), value)
+		else
+			_config = _config.replace(new RegExp(`\\$\\{${key}:([\\w]+)\\}`, "g"), (match, name) => {
+				if (value[name] == null)
+					console.warn(`[injectVariables] variable "${name}" referenced but not found in "${key}"`)
 
-		return process.env[name] ?? notFoundValue
-	})
+				return value[name] ?? propNotFoundValue ?? match
+			})
+	}
 
 	return (isObject ? JSON.parse(_config) : _config) as C extends string ? string : C
 }