Pass metadata context through tool pipeline to MCP calls (#1323)

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/ArtifactSinkingTool.kt

@@ -88,8 +88,11 @@ class ArtifactSinkingTool<T : Any>(
     override val definition: Tool.Definition = delegate.definition
     override val metadata: Tool.Metadata = delegate.metadata
 
-    override fun call(input: String): Tool.Result {
-        val result = delegate.call(input)
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        callAndSink { delegate.call(input, context) }
+
+    private inline fun callAndSink(action: () -> Tool.Result): Tool.Result {
+        val result = action()
 
         if (result is Tool.Result.WithArtifact) {
             val artifact = result.artifact

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/DelegatingTool.kt

@@ -19,11 +19,39 @@ package com.embabel.agent.api.tool
  * Interface for tool decorators that wrap another tool.
  * Enables unwrapping to find the underlying tool implementation.
  * Thus, it is important that tool wrappers implement this interface to allow unwrapping.
+ *
+ * ## Canonical call method
+ *
+ * [call] (String, ToolCallContext) is the **single canonical entry point** for
+ * decorator logic. Decorators should override only this method. The single-arg
+ * [call] (String) routes through it automatically via [ToolCallContext.EMPTY],
+ * so both call paths execute the same decorator behavior.
+ *
+ * This eliminates a class of bugs where a decorator overrides [call] (String)
+ * but the two-arg variant (used by [com.embabel.agent.spi.loop.support.DefaultToolLoop])
+ * bypasses the decorator entirely.
  */
 interface DelegatingTool : Tool {
 
     /**
      * The underlying tool being delegated to.
      */
     val delegate: Tool
+
+    /**
+     * Routes single-arg calls through the canonical two-arg method,
+     * ensuring decorator logic in [call] (String, ToolCallContext) is
+     * always executed regardless of which overload the caller uses.
+     */
+    override fun call(input: String): Tool.Result =
+        call(input, ToolCallContext.EMPTY)
+
+    /**
+     * Canonical entry point for decorator logic. Override this method
+     * to add behavior while preserving context propagation to [delegate].
+     *
+     * The default implementation simply forwards to the delegate.
+     */
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        delegate.call(input, context)
 }

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/MethodTool.kt

@@ -32,6 +32,11 @@ import kotlin.reflect.jvm.javaType
 
 /**
  * Tool implementation that wraps a method annotated with [@LlmTool].
+ *
+ * Supports [ToolCallContext] injection: if the annotated method declares a
+ * parameter of type [ToolCallContext], the framework injects the current context
+ * automatically — just like Spring AI injects `ToolContext` into `@Tool` methods.
+ * Such parameters are excluded from the JSON input schema sent to the LLM.
  */
 internal sealed class MethodTool(
     protected val instance: Any,
@@ -43,10 +48,16 @@ internal sealed class MethodTool(
 
     override val metadata: Tool.Metadata = Tool.Metadata(returnDirect = annotation.returnDirect)
 
-    override fun call(input: String): Tool.Result {
+    override fun call(input: String): Tool.Result =
+        callWithContext(input, ToolCallContext.EMPTY)
+
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        callWithContext(input, context)
+
+    private fun callWithContext(input: String, context: ToolCallContext): Tool.Result {
         return try {
             val args = parseArguments(input)
-            val result = invokeMethod(args)
+            val result = invokeMethod(args, context)
             convertResult(result)
         } catch (e: Exception) {
             // Unwrap InvocationTargetException to get the actual cause
@@ -74,7 +85,7 @@ internal sealed class MethodTool(
         }
     }
 
-    protected abstract fun invokeMethod(args: Map<String, Any?>): Any?
+    protected abstract fun invokeMethod(args: Map<String, Any?>, context: ToolCallContext): Any?
 
     private fun convertResult(result: Any?): Tool.Result {
         return when (result) {
@@ -138,8 +149,10 @@ internal class KotlinMethodTool(
     override val definition: Tool.Definition by lazy {
         val name = annotation.name.ifEmpty { method.name }
         // Use victools-based schema generation for proper generic type handling
+        // Exclude ToolCallContext parameters — they are framework-injected, not LLM-provided
         val parameterInfos = method.parameters
             .filter { it.kind == KParameter.Kind.VALUE }
+            .filter { it.type.javaType != ToolCallContext::class.java }
             .map { param ->
                 val paramAnnotation = param.findAnnotation<Param>()
                 ParameterInfo(
@@ -156,7 +169,7 @@ internal class KotlinMethodTool(
         )
     }
 
-    override fun invokeMethod(args: Map<String, Any?>): Any? {
+    override fun invokeMethod(args: Map<String, Any?>, context: ToolCallContext): Any? {
         val params = method.parameters
         val callArgs = mutableMapOf<KParameter, Any?>()
 
@@ -165,6 +178,12 @@ internal class KotlinMethodTool(
             when (param.kind) {
                 KParameter.Kind.INSTANCE -> callArgs[param] = instance
                 KParameter.Kind.VALUE -> {
+                    // Inject ToolCallContext if the parameter type matches
+                    if (param.type.javaType == ToolCallContext::class.java) {
+                        callArgs[param] = context
+                        continue
+                    }
+
                     val paramAnnotation = param.findAnnotation<Param>()
                     val paramName = param.name ?: continue
                     val value = args[paramName]
@@ -207,7 +226,9 @@ internal class JavaMethodTool(
     override val definition: Tool.Definition by lazy {
         val name = annotation.name.ifEmpty { method.name }
         // Use victools-based schema generation for proper generic type handling
+        // Exclude ToolCallContext parameters — they are framework-injected, not LLM-provided
         val parameterInfos = method.parameters
+            .filter { !ToolCallContext::class.java.isAssignableFrom(it.type) }
             .map { param ->
                 val paramAnnotation = param.getAnnotation(Param::class.java)
                 ParameterInfo(
@@ -224,11 +245,16 @@ internal class JavaMethodTool(
         )
     }
 
-    override fun invokeMethod(args: Map<String, Any?>): Any? {
+    override fun invokeMethod(args: Map<String, Any?>, context: ToolCallContext): Any? {
         val params = method.parameters
         val callArgs = arrayOfNulls<Any?>(method.parameters.size)
 
         for ((index, param) in params.withIndex()) {
+            // Inject ToolCallContext if the parameter type matches
+            if (ToolCallContext::class.java.isAssignableFrom(param.type)) {
+                callArgs[index] = context
+                continue
+            }
             val value = args[param.name]
             if (value != null) {
                 // Convert value to expected type if needed

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/ReplanningTools.kt

@@ -43,9 +43,6 @@ fun interface ReplanningToolBlackboardUpdater {
  * - Chat routing: A routing tool classifies intent and triggers replan to switch handlers
  * - Discovery: A tool discovers information that requires a different plan
  *
- * Note: This tool accesses [AgentProcess] via thread-local at call time, which is set
- * by the decorator chain.
- *
  * @param delegate The tool to wrap
  * @param reason Human-readable explanation of why replan is needed
  * @param blackboardUpdater Callback to update the blackboard before replanning.
@@ -65,8 +62,11 @@ class ReplanningTool @JvmOverloads constructor(
     override val definition: Tool.Definition = delegate.definition
     override val metadata: Tool.Metadata = delegate.metadata
 
-    override fun call(input: String): Tool.Result {
-        val result = delegate.call(input)
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        callAndReplan { delegate.call(input, context) }
+
+    private inline fun callAndReplan(action: () -> Tool.Result): Tool.Result {
+        val result = action()
         val resultContent = result.content
 
         throw ReplanRequestedException(
@@ -135,9 +135,6 @@ fun interface ReplanDecider {
  * Unlike [ReplanningTool] which always triggers replanning, this tool allows the [ReplanDecider]
  * to inspect the result and decide whether to replan.
  *
- * Note: This tool accesses [AgentProcess] via thread-local at call time, which is set
- * by the decorator chain.
- *
  * @param delegate The tool to wrap
  * @param decider Decider that inspects the result context and determines whether to replan
  */
@@ -149,19 +146,22 @@ class ConditionalReplanningTool(
     override val definition: Tool.Definition = delegate.definition
     override val metadata: Tool.Metadata = delegate.metadata
 
-    override fun call(input: String): Tool.Result {
-        val result = delegate.call(input)
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        callAndMaybeReplan { delegate.call(input, context) }
+
+    private inline fun callAndMaybeReplan(action: () -> Tool.Result): Tool.Result {
+        val result = action()
 
         val agentProcess = AgentProcess.get()
             ?: throw IllegalStateException("No AgentProcess available for ConditionalReplanningTool")
 
-        val context = ReplanContext(
+        val replanContext = ReplanContext(
             result = result,
             agentProcess = agentProcess,
             tool = delegate,
         )
 
-        val decision = decider.evaluate(context)
+        val decision = decider.evaluate(replanContext)
         if (decision != null) {
             throw ReplanRequestedException(
                 reason = decision.reason,

embabel-agent-api/src/main/kotlin/com/embabel/agent/api/tool/Tool.kt

@@ -51,6 +51,22 @@ interface Tool : ToolInfo {
      */
     fun call(input: String): Result
 
+    /**
+     * Execute the tool with JSON input and out-of-band context.
+     *
+     * The default implementation simply delegates to [call] (String),
+     * discarding the context. Override this method to receive context
+     * explicitly (e.g., for auth tokens, tenant IDs, or correlation IDs).
+     *
+     * [DelegatingTool] provides a default that propagates context through
+     * decorator chains, so most decorators do not need to override this.
+     *
+     * @param input JSON string matching inputSchema
+     * @param context out-of-band metadata (auth tokens, tenant IDs, etc.)
+     * @return Result to send back to LLM
+     */
+    fun call(input: String, context: ToolCallContext): Result = call(input)
+
     /**
      * Framework-agnostic tool definition.
      */
@@ -259,6 +275,14 @@ interface Tool : ToolInfo {
         fun invoke(input: String): Result
     }
 
+    /**
+     * Functional interface for context-aware tool implementations.
+     * Use when the tool needs out-of-band metadata (auth tokens, tenant IDs, etc.).
+     */
+    fun interface ContextAwareFunction {
+        fun invoke(input: String, context: ToolCallContext): Result
+    }
+
     /**
      * Java-friendly functional interface for tool implementations.
      * Uses `handle` method name which is more idiomatic in Java than `invoke`.
@@ -295,6 +319,32 @@ interface Tool : ToolInfo {
             function: Function,
         ): Tool = of(name, description, InputSchema.empty(), metadata, function)
 
+        /**
+         * Create a context-aware tool from a [ContextAwareFunction].
+         * The function receives [ToolCallContext] explicitly at call time.
+         */
+        fun of(
+            name: String,
+            description: String,
+            inputSchema: InputSchema,
+            metadata: Metadata = Metadata.DEFAULT,
+            function: ContextAwareFunction,
+        ): Tool = ContextAwareFunctionalTool(
+            definition = Definition(name, description, inputSchema),
+            metadata = metadata,
+            function = function,
+        )
+
+        /**
+         * Create a context-aware tool with no parameters.
+         */
+        fun of(
+            name: String,
+            description: String,
+            metadata: Metadata = Metadata.DEFAULT,
+            function: ContextAwareFunction,
+        ): Tool = of(name, description, InputSchema.empty(), metadata, function)
+
         /**
          * Create a tool with no parameters (Java-friendly).
          * This method is easier to call from Java as it uses the Handler interface.
@@ -572,8 +622,6 @@ private class RenamedTool(
 
     override val metadata: Tool.Metadata
         get() = delegate.metadata
-
-    override fun call(input: String): Tool.Result = delegate.call(input)
 }
 
 /**
@@ -593,8 +641,6 @@ private class DescribedTool(
 
     override val metadata: Tool.Metadata
         get() = delegate.metadata
-
-    override fun call(input: String): Tool.Result = delegate.call(input)
 }
 
 // Private implementations
@@ -698,3 +744,15 @@ private class FunctionalTool(
     override fun call(input: String): Tool.Result =
         function.invoke(input)
 }
+
+private class ContextAwareFunctionalTool(
+    override val definition: Tool.Definition,
+    override val metadata: Tool.Metadata,
+    private val function: Tool.ContextAwareFunction,
+) : Tool {
+    override fun call(input: String): Tool.Result =
+        function.invoke(input, ToolCallContext.EMPTY)
+
+    override fun call(input: String, context: ToolCallContext): Tool.Result =
+        function.invoke(input, context)
+}