spring-projects / spring-projects/spring-ai
Auth (and others) context lost on MCP Tool call when using streaming chat model
Nobody has claimed this yet.
- 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
AuthToolCallingManagerwith 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
AuthToolCallingManageraround 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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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