spring-projects / spring-projects/spring-ai

Auth (and others) context lost on MCP Tool call when using streaming chat model

Open
#3,877 1 comment 2 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

status: waiting-for-triage
Dominant language
Java
Stars
9.5k
Forks
2.9k
Avg merge
1d 7h
Merged PRs (30d)
6

Description

Bug description
All Request context (RequestAttribute, Security) missing on tool calling when using streaming chat model with stream response, which is necessary on MCP tool call.

Following this tutorial on tool calling it will check based on RequestContextHolder as well.

    if (RequestContextHolder.getRequestAttributes() instanceof ServletRequestAttributes) {
      return this.delegate.filter(request, next);
    }

I am using streaming AzureOpenAiChatModel, this is the scenario:

  • Request come -> Reactor thread
  • Processing prompt etc -> BoundedElastic -> Context still can be preserved using scheduledHook
  • AzureChat call model -> ReactorThread -> Context lost, cant use scheduledHook since its not on scheduled thread anymore
  • CallTool -> Context does not exist, and if we follow above tutorial, it will use client credentials, which is incorrect

Environment

  • Spring AI 1.0.0
  • Java 21

Steps to reproduce

  • Use auth_code flow
  • Use streaming-chat-model (I use AzureAI, not sure if the issue occurs at other models)
  • Make use of any remote MCP Server
  • Follow Spring tutorial for MCP Auth
  • User auth_code flow context will be lost on the tool calling

Expected behavior

  • SecurityContext should be retained (whether we are using webmvc or webflux), currently either way does not work
  • Even if it is lost, at least we should be able to propagate it using scheduledWebhook, all the subscribers on need to be on the boundedElastic thread

NB - scheduledWebhook example that I also used at MCP SDK Server context issue:

    @PostConstruct
    public void init() {
        // This is the key for propagating context to Reactor's boundedElastic scheduler
        Function<Runnable, Runnable> decorator = runnable -> {
            // This part is executed on the *submitting* thread (e.g. web request thread)
            Authentication auth = SecurityContextHolder.getContext().getAuthentication();
            RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
            return () -> {
                try {
                    // This part is executed on the *worker* thread (e.g. boundedElastic-1)
                    if (SecurityContextHolder.getContext().getAuthentication() != null) {
                        log.warn(
                                "SecurityContext is already exist: {}, overwriting it...",
                                SecurityContextHolder.getContext().getAuthentication());
                    }
                    SecurityContextHolder.getContext().setAuthentication(auth);
                    RequestContextHolder.setRequestAttributes(requestAttributes);
                    runnable.run();
                } finally {
                    SecurityContextHolder.clearContext();
                    RequestContextHolder.resetRequestAttributes();
                }
            };
        };

        Schedulers.onScheduleHook("McpBoundedElasticHook", decorator);
    }

Workaround

  • Create custom AuthToolCallingManager with ObjectProvider that will delegate to the default/custom authToolCallingManager, store the context when it is still available
@Scope("prototype")
@RequiredArgsConstructor
public class AuthToolCallingManager implements ToolCallingManager {

    private Authentication currentAuth = SecurityContextHolder.getContext().getAuthentication();
    private RequestAttributes currentRequestAttributes = RequestContextHolder.getRequestAttributes();

    private final ToolCallingManager delegate;

    @Override
    public List<ToolDefinition> resolveToolDefinitions(ToolCallingChatOptions chatOptions) {
        return delegate.resolveToolDefinitions(chatOptions);
    }

    @Override
    public ToolExecutionResult executeToolCalls(Prompt prompt, ChatResponse chatResponse) {
        SecurityContextHolder.getContext().setAuthentication(currentAuth);
        RequestContextHolder.setRequestAttributes(currentRequestAttributes);
        try {
            return delegate.executeToolCalls(prompt, chatResponse);
        } finally {
            SecurityContextHolder.clearContext();
            RequestContextHolder.resetRequestAttributes();
        }
    }

    public void clear() {
        currentAuth = null;
        currentRequestAttributes = null;
    }
}
  • Use manual tool calling with .internalToolExecutionEnabled(false)
  • Pass the AuthToolCallingManager around on your manual tool execution, something like below:

    private Flux<ChatResponse> processResponse(
            String conversationId, ChatResponse response, AuthToolCallingManager authTcm) {
        if (response.hasToolCalls()) {
            // Handle tool calls manually and continue conversation
            return handleToolCallsAndContinue(conversationId, response, authTcm);
        } else {
            return Flux.just(response);
        }
    }

    private Flux<ChatResponse> handleToolCallsAndContinue(
            String conversationId, ChatResponse toolCallIntention, AuthToolCallingManager authTcm) {
        // Store tool intention to memory
        chatMemory.add(conversationId, toolCallIntention.getResult().getOutput());
        ToolExecutionResult toolExecutionResult = authTcm.executeToolCalls(
                promptFromMemory(conversationId), // Prompt came from memory
                toolCallIntention);
        Message toolResultMessage = toolExecutionResult
                .conversationHistory()
                .get(toolExecutionResult.conversationHistory().size() - 1);

        // Store tool result to memory 
        chatMemory.add(conversationId, toolResultMessage);

        // Continue conversation with tool results
        // This will make a duplicate of user message if using chatMemory refer to: https://github.com/spring-projects/spring-ai/issues/2101#issuecomment-3101687749
        return buildRequestSpec(promptFromMemory(conversationId), conversationId, false).stream()
                .chatResponse()
                .flatMap(response -> processResponse(conversationId, response, authTcm))
                .startWith(toolCallIntention); // Include the original prompt
    }

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Reproduce the streaming scenario with AzureOpenAiChatModel, the MCP OAuth2 tutorial, and a remote MCP server, then trace the flow from the streaming model into ToolCallingManager.executeToolCalls. Done means SecurityContext and RequestAttributes remain available during tool calls for both WebMVC and WebFlux, or can be propagated as described.

Written by the indexing model from the issue text.

Assessment

Tech stack
azure, java, spring
Domain
api, authentication, backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.