@@ -8,6 +8,7 @@ use crossterm::{
88 style,
99} ;
1010use eyre:: Result ;
11+ use regex:: Regex ;
1112use schemars:: JsonSchema ;
1213use serde:: {
1314 Deserialize ,
@@ -23,6 +24,7 @@ use crate::cli::agent::{
2324} ;
2425use crate :: cli:: chat:: CONTINUATION_LINE ;
2526use crate :: cli:: chat:: token_counter:: TokenCounter ;
27+ use crate :: os:: Os ;
2628use crate :: mcp_client:: {
2729 Client as McpClient ,
2830 ClientConfig as McpClientConfig ,
@@ -35,7 +37,6 @@ use crate::mcp_client::{
3537 StdioTransport ,
3638 ToolCallResult ,
3739} ;
38- use crate :: os:: Os ;
3940
4041// TODO: support http transport type
4142#[ derive( Clone , Serialize , Deserialize , Debug , Eq , PartialEq , JsonSchema ) ]
@@ -63,6 +64,26 @@ pub fn default_timeout() -> u64 {
6364 120 * 1000
6465}
6566
67+ /// Substitutes environment variables in the format ${env:VAR_NAME} with their actual values
68+ fn substitute_env_vars ( input : & str , env : & crate :: os:: Env ) -> String {
69+ // Create a regex to match ${env:VAR_NAME} pattern
70+ let re = Regex :: new ( r"\$\{env:([^}]+)\}" ) . unwrap ( ) ;
71+
72+ re. replace_all ( input, |caps : & regex:: Captures < ' _ > | {
73+ let var_name = & caps[ 1 ] ;
74+ env. get ( var_name) . unwrap_or_else ( |_| format ! ( "${{{}}}" , var_name) )
75+ } )
76+ . to_string ( )
77+ }
78+
79+ /// Process a HashMap of environment variables, substituting any ${env:VAR_NAME} patterns
80+ /// with their actual values from the environment
81+ fn process_env_vars ( env_vars : & mut HashMap < String , String > , env : & crate :: os:: Env ) {
82+ for ( _, value) in env_vars. iter_mut ( ) {
83+ * value = substitute_env_vars ( value, env) ;
84+ }
85+ }
86+
6687#[ derive( Debug ) ]
6788pub enum CustomToolClient {
6889 Stdio {
@@ -75,7 +96,7 @@ pub enum CustomToolClient {
7596
7697impl CustomToolClient {
7798 // TODO: add support for http transport
78- pub fn from_config ( server_name : String , config : CustomToolConfig ) -> Result < Self > {
99+ pub fn from_config ( server_name : String , config : CustomToolConfig , os : & crate :: os :: Os ) -> Result < Self > {
79100 let CustomToolConfig {
80101 command,
81102 args,
@@ -84,6 +105,13 @@ impl CustomToolClient {
84105 disabled : _,
85106 ..
86107 } = config;
108+
109+ // Process environment variables if present
110+ let processed_env = env. map ( |mut env_vars| {
111+ process_env_vars ( & mut env_vars, & os. env ) ;
112+ env_vars
113+ } ) ;
114+
87115 let mcp_client_config = McpClientConfig {
88116 server_name : server_name. clone ( ) ,
89117 bin_path : command. clone ( ) ,
@@ -93,7 +121,7 @@ impl CustomToolClient {
93121 "name" : "Q CLI Chat" ,
94122 "version" : "1.0.0"
95123 } ) ,
96- env,
124+ env : processed_env ,
97125 } ;
98126 let client = McpClient :: < JsonRpcStdioTransport > :: from_config ( mcp_client_config) ?;
99127 Ok ( CustomToolClient :: Stdio {
@@ -279,3 +307,52 @@ impl CustomTool {
279307 }
280308 }
281309}
310+
311+ #[ cfg( test) ]
312+ mod tests {
313+ use super :: * ;
314+
315+ #[ tokio:: test]
316+ async fn test_substitute_env_vars ( ) {
317+ // Set a test environment variable
318+ let os = Os :: new ( ) . await . unwrap ( ) ;
319+ unsafe {
320+ os. env . set_var ( "TEST_VAR" , "test_value" ) ;
321+ }
322+
323+ // Test basic substitution
324+ assert_eq ! ( substitute_env_vars( "Value is ${env:TEST_VAR}" , & os. env) , "Value is test_value" ) ;
325+
326+ // Test multiple substitutions
327+ assert_eq ! (
328+ substitute_env_vars( "${env:TEST_VAR} and ${env:TEST_VAR}" , & os. env) ,
329+ "test_value and test_value"
330+ ) ;
331+
332+ // Test non-existent variable
333+ assert_eq ! ( substitute_env_vars( "${env:NON_EXISTENT_VAR}" , & os. env) , "${NON_EXISTENT_VAR}" ) ;
334+
335+ // Test mixed content
336+ assert_eq ! (
337+ substitute_env_vars( "Prefix ${env:TEST_VAR} suffix" , & os. env) ,
338+ "Prefix test_value suffix"
339+ ) ;
340+ }
341+
342+ #[ tokio:: test]
343+ async fn test_process_env_vars ( ) {
344+ let os = Os :: new ( ) . await . unwrap ( ) ;
345+ unsafe {
346+ os. env . set_var ( "TEST_VAR" , "test_value" ) ;
347+ }
348+
349+ let mut env_vars = HashMap :: new ( ) ;
350+ env_vars. insert ( "KEY1" . to_string ( ) , "Value is ${env:TEST_VAR}" . to_string ( ) ) ;
351+ env_vars. insert ( "KEY2" . to_string ( ) , "No substitution" . to_string ( ) ) ;
352+
353+ process_env_vars ( & mut env_vars, & os. env ) ;
354+
355+ assert_eq ! ( env_vars. get( "KEY1" ) . unwrap( ) , "Value is test_value" ) ;
356+ assert_eq ! ( env_vars. get( "KEY2" ) . unwrap( ) , "No substitution" ) ;
357+ }
358+ }
0 commit comments