ToolContext with MCP

[rafaelrddc]

Please do a quick search on GitHub issues first, the feature you are about to request might have already been requested.

Expected Behavior

Enable the use of ToolContext together with MCP.

Current Behavior

I looked in the documentation but I didn't find anything referring to how to use ToolContext together with MCP.

Context

I would like it to be possible to use the values ​​available in the ToolContext to call tools provided by MCP servers.

[henningSaulCM]

Related to https://github.com/spring-projects/spring-ai/issues/2378

[He-Pin]

I think we need pass the ServerRequest to the ToolContext

currently, only the body is passed.

[rafaelrddc]

Hello everyone!

Since the MCP specification does not yet provide for this use case, I was able to solve the problem by creating an Advisor that adds the context to the prompt. And so far it has worked well.

class ToolContextPromptAdvisor : AbstractAdvisor(Ordered.HIGHEST_PRECEDENCE) {
    private companion object {
        const val TOOL_CONTEXT_KEY = "tool_context"

        const val PROMPT_TOOL_CONTEXT =
            """
                # TOOL CONTEXT:
                Below is the tool context data provided for your use. These values are available to support your task 
                but should NOT be included or mentioned in your responses:

                {$TOOL_CONTEXT_KEY}  
            """
    }

    override fun aroundCall(
        advisedRequest: AdvisedRequest,
        chain: CallAroundAdvisorChain,
    ): AdvisedResponse = chain.nextAroundCall(addToolContextInPrompt(advisedRequest))

    private fun addToolContextInPrompt(advisedRequest: AdvisedRequest): AdvisedRequest {
        val toolContext = advisedRequest.toolContext
            .map { "- *${it.key}*: ${it.value}" }
            .joinToString(separator = System.lineSeparator())

        val systemParams = advisedRequest.systemParams.toMutableMap()
        systemParams[TOOL_CONTEXT_KEY] = toolContext
        systemParams[CURRENT_DATE_TIME] = LocalDateTime.now().toString()

        return AdvisedRequest
            .from(advisedRequest)
            .toolContext(emptyMap())
            .systemParams(systemParams)
            .systemText(advisedRequest.systemText + System.lineSeparator() + PROMPT_TOOL_CONTEXT)
            .build()
    }
}

[cidd04]

@rafaelrddc This is great. Is it prone to prompt injection though? Have you tested prompts that could overwrite values in your toolcontext?

[rafaelrddc]

@cidd04 This approach is indeed prone to prompt injection, as the user can send this information in the message and the AI ​​uses it instead of the data added by the Advisor.

[jeweis]

I am also facing this issue, is there any alternative solution?

[rafaelrddc]

Hello!

I'm testing a second way to pass context information to the function, which I believe removes the prompt inject problem. I am doing it as follows. On the server side I declare a parameter as optional. Example

Server

@Tool(description = "Set a user alarm for the given time, provided in ISO-8601 format")
void setAlarm(String time, @ToolParam(required = false) UUID userId) {
   if (Objects.isNull(userId)) {
       return "Failed to process the request." // Generic error for AI not informing the user what it needs
   }

   LocalDateTime alarmTime = LocalDateTime.parse(time, DateTimeFormatter.ISO_DATE_TIME);
   System.out.println("Alarm set for " + alarmTime);
}

On the client side, the ToolCallback class was implemented where I add context information to the call to the MCP.

Client

class AsyncMcpToolCallback(
    private val mcpAsyncClient: McpAsyncClient,
    private val tool: McpSchema.Tool,
): ToolCallback {
    override fun call(toolInput: String): String {
        return call(toolInput, null)
    }

    override fun call(toolInput: String, tooContext: ToolContext?): String {
        val arguments = ModelOptionsUtils.jsonToMap(toolInput) +
                (tooContext?.context ?: emptyMap())

        return mcpAsyncClient.callTool(McpSchema.CallToolRequest(tool.name, arguments))
            .map {
                ModelOptionsUtils.toJsonString(it.content)
            }.block().toString()
    }

    override fun getToolDefinition(): ToolDefinition {
        return ToolDefinition.builder().name(
            McpToolUtils.prefixedToolName(
                this.mcpAsyncClient.clientInfo.name(),
                tool.name()
            )
        ).description(tool.description()).inputSchema(
            ModelOptionsUtils.toJsonString(
                tool.inputSchema()
            )
        ).build()
    }
}