diff --git a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
index 0c34e0644..9d13c6953 100644
--- a/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
+++ b/src/main/java/org/opensearch/flowframework/workflow/ToolStep.java
@@ -88,6 +88,7 @@ public PlainActionFuture<WorkflowData> execute(
                 outputs,
                 toolParameterKeys
             );
+            @SuppressWarnings("unchecked")
             Map<String, String> config = (Map<String, String>) inputs.getOrDefault(CONFIG_FIELD, Collections.emptyMap());
 
             MLToolSpec.MLToolSpecBuilder builder = MLToolSpec.builder();
@@ -105,8 +106,7 @@ public PlainActionFuture<WorkflowData> execute(
             if (includeOutputInAgentResponse != null) {
                 builder.includeOutputInAgentResponse(includeOutputInAgentResponse);
             }
-            // TODO https://github.com/opensearch-project/ml-commons/pull/2977/files must be merged for this to compile
-            // builder.configMap(config);
+            builder.configMap(config);
 
             MLToolSpec mlToolSpec = builder.build();
 
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
index c2b3dcca1..8def95f58 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
@@ -56,7 +56,13 @@ public void setUp() throws Exception {
         this.flowFrameworkIndicesHandler = mock(FlowFrameworkIndicesHandler.class);
         MockitoAnnotations.openMocks(this);
 
-        MLToolSpec tools = new MLToolSpec("tool1", "CatIndexTool", "desc", Collections.emptyMap(), false);
+        MLToolSpec tools = MLToolSpec.builder()
+            .type("tool1")
+            .name("CatIndexTool")
+            .description("desc")
+            .parameters(Collections.emptyMap())
+            .includeOutputInAgentResponse(false)
+            .build();
 
         LLMSpec llmSpec = new LLMSpec("xyz", Collections.emptyMap());
 
diff --git a/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java b/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java
index 0ef0e271a..2b5e5b7fa 100644
--- a/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java
+++ b/src/test/java/org/opensearch/flowframework/workflow/ToolStepTests.java
@@ -103,8 +103,7 @@ public void testTool() throws ExecutionException, InterruptedException {
         );
         assertTrue(future.isDone());
         assertEquals(MLToolSpec.class, future.get().getContent().get("tools").getClass());
-        // TODO https://github.com/opensearch-project/ml-commons/pull/2977/files must be merged for this to compile
-        // assertEquals(Map.of("foo", "bar"), ((MLToolSpec) future.get().getContent().get("tools")).getConfigMap());
+        assertEquals(Map.of("foo", "bar"), ((MLToolSpec) future.get().getContent().get("tools")).getConfigMap());
     }
 
     public void testBoolParseFail() {