fix(ai): propagate security context to tool callbacks

Signed-off-by: yuluo-yx <yuluo08290126@gmail.com>

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/McpContextHolder.java

@@ -18,14 +18,21 @@
 package org.apache.hertzbeat.ai.config;
 
 import com.usthe.sureness.subject.SubjectSum;
-import org.springframework.core.NamedInheritableThreadLocal;
+import com.usthe.sureness.util.SurenessContextHolder;
+import java.util.Map;
+import java.util.function.Supplier;
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.core.NamedThreadLocal;
 
 /**
  * Context holder for AI agent security context.
  */
 public final class McpContextHolder {
+
+    static final String SUBJECT_CONTEXT_KEY = McpContextHolder.class.getName() + ".subject";
+
     private static final ThreadLocal<SubjectSum> subjectHolder =
-            new NamedInheritableThreadLocal<>("MCP Security and User Identification Context");
+            new NamedThreadLocal<>("MCP Security and User Identification Context");
 
     private McpContextHolder() {}
 
@@ -44,10 +51,73 @@ public final class McpContextHolder {
         return subjectHolder.get();
     }
 
+    /**
+     * Creates a security context that is propagated only through the Spring AI tool invocation chain.
+     *
+     * @param subject current authenticated subject, may be {@code null}
+     * @return tool context without null values
+     */
+    public static Map<String, Object> createToolContext(SubjectSum subject) {
+        return subject == null ? Map.of() : Map.of(SUBJECT_CONTEXT_KEY, subject);
+    }
+
+    /**
+     * Retrieves and validates the authenticated subject from the Spring AI tool context.
+     *
+     * @param toolContext Spring AI tool context
+     * @return authenticated subject, or {@code null} when absent
+     */
+    public static SubjectSum getSubject(ToolContext toolContext) {
+        if (toolContext == null) {
+            return null;
+        }
+        Object subject = toolContext.getContext().get(SUBJECT_CONTEXT_KEY);
+        return subject instanceof SubjectSum subjectSum ? subjectSum : null;
+    }
+
+    /**
+     * Executes a synchronous tool invocation within the given subject scope and restores the original thread
+     * contexts afterward.
+     *
+     * <p>Spring AI tool callbacks may run on pooled threads, so they cannot rely on a {@link ThreadLocal} left on
+     * the Servlet request thread. This method binds both MCP and Sureness contexts so the tool and downstream
+     * services observe the same user identity.</p>
+     *
+     * @param subject current request subject, may be {@code null}
+     * @param operation synchronous tool invocation
+     * @param <T> tool result type
+     * @return tool invocation result
+     */
+    public static <T> T callWithSubject(SubjectSum subject, Supplier<T> operation) {
+        SubjectSum previousMcpSubject = getSubject();
+        SubjectSum previousSurenessSubject = SurenessContextHolder.getBindSubject();
+        replaceSubjects(subject);
+        try {
+            return operation.get();
+        } finally {
+            replaceSubjects(previousMcpSubject, previousSurenessSubject);
+        }
+    }
+
     /**
      * Clears the context from the thread to prevent memory leaks.
      */
     public static void clear() {
         subjectHolder.remove();
     }
+
+    private static void replaceSubjects(SubjectSum subject) {
+        replaceSubjects(subject, subject);
+    }
+
+    private static void replaceSubjects(SubjectSum mcpSubject, SubjectSum surenessSubject) {
+        clear();
+        SurenessContextHolder.unbindSubject();
+        if (mcpSubject != null) {
+            setSubject(mcpSubject);
+        }
+        if (surenessSubject != null) {
+            SurenessContextHolder.bindSubject(surenessSubject);
+        }
+    }
 }

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/SecurityContextToolCallback.java

@@ -0,0 +1,60 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.ai.config;
+
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.ai.tool.ToolCallback;
+import org.springframework.ai.tool.definition.ToolDefinition;
+import org.springframework.ai.tool.metadata.ToolMetadata;
+
+/**
+ * Establishes a request-scoped security context for Spring AI tool callbacks.
+ *
+ * <p>The wrapper keeps the tool definition and metadata unchanged. It binds the identity only for the delegated
+ * call and restores the worker thread's original state after both successful and failed invocations.</p>
+ */
+public final class SecurityContextToolCallback implements ToolCallback {
+
+    private final ToolCallback delegate;
+
+    public SecurityContextToolCallback(ToolCallback delegate) {
+        this.delegate = delegate;
+    }
+
+    @Override
+    public ToolDefinition getToolDefinition() {
+        return delegate.getToolDefinition();
+    }
+
+    @Override
+    public ToolMetadata getToolMetadata() {
+        return delegate.getToolMetadata();
+    }
+
+    @Override
+    public String call(String toolInput) {
+        return call(toolInput, null);
+    }
+
+    @Override
+    public String call(String toolInput, ToolContext toolContext) {
+        return McpContextHolder.callWithSubject(
+                McpContextHolder.getSubject(toolContext),
+                () -> delegate.call(toolInput, toolContext));
+    }
+}

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/controller/ChatController.java

@@ -17,13 +17,10 @@
 
 package org.apache.hertzbeat.ai.controller;
 
-import com.usthe.sureness.subject.SubjectSum;
-import com.usthe.sureness.util.SurenessContextHolder;
 import io.swagger.v3.oas.annotations.Operation;
 import io.swagger.v3.oas.annotations.Parameter;
 import io.swagger.v3.oas.annotations.tags.Tag;
 import lombok.extern.slf4j.Slf4j;
-import org.apache.hertzbeat.ai.config.McpContextHolder;
 import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
 import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
 import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
@@ -75,8 +72,6 @@ public class ChatController {
     public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(@Valid @RequestBody ChatRequestContext context) {
         try {
             // Validate message is not empty
-            SubjectSum subject = SurenessContextHolder.getBindSubject();
-            McpContextHolder.setSubject(subject);
             if (context.getMessage() == null || context.getMessage().trim().isEmpty()) {
                 ChatResponseChunk errorResponse = ChatResponseChunk.builder()
                     .conversationId(context.getConversationId())

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/pojo/dto/ChatRequestContext.java

@@ -18,12 +18,13 @@
 
 package org.apache.hertzbeat.ai.pojo.dto;
 
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.usthe.sureness.subject.SubjectSum;
+import java.util.List;
 import lombok.AllArgsConstructor;
 import lombok.Builder;
 import lombok.Data;
 import lombok.NoArgsConstructor;
-
-import java.util.List;
 import org.apache.hertzbeat.common.entity.ai.ChatMessage;
 
 /**
@@ -47,4 +48,11 @@ public class ChatRequestContext {
      * Conversation history messages for context
      */
     private List<ChatMessage> conversationHistory;
+
+    /**
+     * Authenticated subject captured by the server. It cannot be supplied through client JSON and is never included
+     * in model messages.
+     */
+    @JsonIgnore
+    private SubjectSum subject;
 }

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImpl.java

@@ -19,10 +19,13 @@
 package org.apache.hertzbeat.ai.service.impl;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.hertzbeat.ai.config.McpContextHolder;
+import org.apache.hertzbeat.ai.config.SecurityContextToolCallback;
 import org.apache.hertzbeat.ai.sop.model.SopDefinition;
 import org.apache.hertzbeat.ai.sop.model.SopParameter;
 import org.apache.hertzbeat.ai.sop.registry.SkillRegistry;
@@ -45,6 +48,7 @@ import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.Message;
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.ai.tool.ToolCallbackProvider;
+import org.springframework.ai.tool.ToolCallback;
 import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.beans.factory.annotation.Qualifier;
 import org.springframework.context.ApplicationContext;
@@ -125,11 +129,15 @@ public class ChatClientProviderServiceImpl implements ChatClientProviderService
             // Build system prompt with dynamic skills list and conversation ID
             // The conversationId is injected into the prompt so AI can pass it to schedule tools
             String systemPrompt = buildSystemPrompt(context.getConversationId());
+            ToolCallback[] toolCallbacks = Arrays.stream(toolCallbackProvider.getToolCallbacks())
+                .map(SecurityContextToolCallback::new)
+                .toArray(ToolCallback[]::new);
 
             return chatClient.prompt()
                 .messages(messages)
                 .system(systemPrompt)
-                .tools(toolCallbackProvider)
+                .tools((Object[]) toolCallbacks)
+                .toolContext(McpContextHolder.createToolContext(context.getSubject()))
                 .stream()
                 .content()
                 .doOnComplete(() -> log.info("Streaming completed for conversation: {}", context.getConversationId()))

hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java

@@ -17,6 +17,8 @@
 
 package org.apache.hertzbeat.ai.service.impl;
 
+import com.usthe.sureness.subject.SubjectSum;
+import com.usthe.sureness.util.SurenessContextHolder;
 import java.util.Optional;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.hertzbeat.ai.dao.ChatConversationDao;
@@ -104,6 +106,8 @@ public class ConversationServiceImpl implements ConversationService {
         // Stream response from AI service
         StringBuilder fullResponse = new StringBuilder();
         ChatMessage finalChatMessage = chatMessage;
+        SubjectSum subject = SurenessContextHolder.getBindSubject();
+        context.setSubject(subject);
         return chatClientProviderService.streamChat(context)
             .map(chunk -> {
                 fullResponse.append(chunk);

hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/config/McpContextHolderTest.java

@@ -0,0 +1,146 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.hertzbeat.ai.config;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.mockito.Mockito.mock;
+
+import com.usthe.sureness.subject.SubjectSum;
+import com.usthe.sureness.util.SurenessContextHolder;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+import org.springframework.ai.chat.model.ToolContext;
+import org.springframework.ai.tool.ToolCallback;
+import org.springframework.ai.tool.definition.ToolDefinition;
+
+/**
+ * Verifies the security context scope used when model tools execute across threads.
+ */
+class McpContextHolderTest {
+
+    @AfterEach
+    void clearContext() {
+        McpContextHolder.clear();
+        SurenessContextHolder.clear();
+    }
+
+    @Test
+    void callWithSubjectShouldBindAndClearBothContexts() {
+        SubjectSum subject = mock(SubjectSum.class);
+
+        String result = McpContextHolder.callWithSubject(subject, () -> {
+            assertSame(subject, McpContextHolder.getSubject());
+            assertSame(subject, SurenessContextHolder.getBindSubject());
+            return "result";
+        });
+
+        assertEquals("result", result);
+        assertNull(McpContextHolder.getSubject());
+        assertNull(SurenessContextHolder.getBindSubject());
+    }
+
+    @Test
+    void callWithSubjectShouldRestoreIndependentPreviousContexts() {
+        SubjectSum previousMcpSubject = mock(SubjectSum.class);
+        SubjectSum previousSurenessSubject = mock(SubjectSum.class);
+        SubjectSum currentSubject = mock(SubjectSum.class);
+        McpContextHolder.setSubject(previousMcpSubject);
+        SurenessContextHolder.bindSubject(previousSurenessSubject);
+
+        McpContextHolder.callWithSubject(currentSubject, () -> {
+            assertSame(currentSubject, McpContextHolder.getSubject());
+            assertSame(currentSubject, SurenessContextHolder.getBindSubject());
+            return null;
+        });
+
+        assertSame(previousMcpSubject, McpContextHolder.getSubject());
+        assertSame(previousSurenessSubject, SurenessContextHolder.getBindSubject());
+    }
+
+    @Test
+    void callWithSubjectShouldRestoreContextsAfterFailure() {
+        SubjectSum previousSubject = mock(SubjectSum.class);
+        SubjectSum currentSubject = mock(SubjectSum.class);
+        McpContextHolder.setSubject(previousSubject);
+        SurenessContextHolder.bindSubject(previousSubject);
+
+        assertThrows(IllegalStateException.class, () ->
+            McpContextHolder.callWithSubject(currentSubject, () -> {
+                throw new IllegalStateException("tool failed");
+            }));
+
+        assertSame(previousSubject, McpContextHolder.getSubject());
+        assertSame(previousSubject, SurenessContextHolder.getBindSubject());
+    }
+
+    @Test
+    void toolContextShouldCarryOnlyValidSubject() {
+        SubjectSum subject = mock(SubjectSum.class);
+        ToolContext toolContext = new ToolContext(McpContextHolder.createToolContext(subject));
+
+        assertSame(subject, McpContextHolder.getSubject(toolContext));
+        assertNull(McpContextHolder.getSubject(null));
+        assertNull(McpContextHolder.getSubject(new ToolContext(
+                java.util.Map.of(McpContextHolder.SUBJECT_CONTEXT_KEY, "invalid"))));
+        assertEquals(java.util.Map.of(), McpContextHolder.createToolContext(null));
+    }
+
+    @Test
+    void callbackShouldExposeSubjectOnlyDuringDelegateCall() {
+        SubjectSum subject = mock(SubjectSum.class);
+        ToolContext toolContext = new ToolContext(McpContextHolder.createToolContext(subject));
+        ToolCallback delegate = new ToolCallback() {
+            @Override
+            public ToolDefinition getToolDefinition() {
+                return ToolDefinition.builder()
+                        .name("test")
+                        .description("test")
+                        .inputSchema("{}")
+                        .build();
+            }
+
+            @Override
+            public String call(String input) {
+                return input;
+            }
+
+            @Override
+            public String call(String input, ToolContext context) {
+                if (context == null) {
+                    return input;
+                }
+                assertSame(subject, McpContextHolder.getSubject());
+                assertSame(subject, SurenessContextHolder.getBindSubject());
+                return "ok";
+            }
+        };
+        SecurityContextToolCallback callback = new SecurityContextToolCallback(delegate);
+
+        assertEquals("test", callback.getToolDefinition().name());
+        assertEquals(delegate.getToolMetadata().returnDirect(), callback.getToolMetadata().returnDirect());
+        assertEquals("plain", callback.call("plain"));
+        String result = callback.call("{}", toolContext);
+
+        assertEquals("ok", result);
+        assertNull(McpContextHolder.getSubject());
+        assertNull(SurenessContextHolder.getBindSubject());
+    }
+}

hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImplTest.java

@@ -18,17 +18,41 @@
 package org.apache.hertzbeat.ai.service.impl;
 
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertInstanceOf;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyList;
+import static org.mockito.ArgumentMatchers.anyMap;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
+import com.usthe.sureness.subject.SubjectSum;
 import java.lang.reflect.Proxy;
+import java.nio.charset.StandardCharsets;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicReference;
+import org.apache.hertzbeat.ai.config.McpContextHolder;
+import org.apache.hertzbeat.ai.config.SecurityContextToolCallback;
+import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
+import org.apache.hertzbeat.ai.sop.registry.SkillRegistry;
 import org.apache.hertzbeat.base.dao.GeneralConfigDao;
 import org.apache.hertzbeat.common.entity.dto.ModelProviderConfig;
 import org.apache.hertzbeat.common.entity.manager.GeneralConfig;
 import org.apache.hertzbeat.common.support.event.AiProviderConfigChangeEvent;
 import org.apache.hertzbeat.common.util.JsonUtil;
 import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.springframework.ai.chat.client.ChatClient;
+import org.springframework.ai.tool.ToolCallback;
+import org.springframework.ai.tool.ToolCallbackProvider;
+import org.springframework.context.ApplicationContext;
 import org.springframework.context.support.StaticApplicationContext;
+import org.springframework.core.io.ByteArrayResource;
+import org.springframework.test.util.ReflectionTestUtils;
+import reactor.core.publisher.Flux;
 
 /**
  * Verifies that the provider configuration cache reacts to enable and disable events.
@@ -52,6 +76,49 @@ class ChatClientProviderServiceImplTest {
         assertFalse(service.isConfigured());
     }
 
+    @Test
+    @SuppressWarnings("unchecked")
+    void streamChatShouldAttachSubjectAndWrapEveryToolCallback() {
+        ApplicationContext applicationContext = mock(ApplicationContext.class);
+        SkillRegistry skillRegistry = mock(SkillRegistry.class);
+        ChatClient chatClient = mock(ChatClient.class);
+        ChatClient.ChatClientRequestSpec requestSpec = mock(ChatClient.ChatClientRequestSpec.class);
+        ChatClient.StreamResponseSpec streamSpec = mock(ChatClient.StreamResponseSpec.class);
+        ToolCallback delegate = mock(ToolCallback.class);
+        SubjectSum subject = mock(SubjectSum.class);
+        ChatClientProviderServiceImpl service = new ChatClientProviderServiceImpl(
+                applicationContext, configDao(new AtomicReference<>()), skillRegistry);
+
+        when(applicationContext.getBean("openAiChatClient", ChatClient.class)).thenReturn(chatClient);
+        when(chatClient.prompt()).thenReturn(requestSpec);
+        when(requestSpec.messages(anyList())).thenReturn(requestSpec);
+        when(requestSpec.system(anyString())).thenReturn(requestSpec);
+        when(requestSpec.tools(any(Object[].class))).thenReturn(requestSpec);
+        when(requestSpec.toolContext(anyMap())).thenReturn(requestSpec);
+        when(requestSpec.stream()).thenReturn(streamSpec);
+        when(streamSpec.content()).thenReturn(Flux.just("answer"));
+        when(skillRegistry.getAllSkills()).thenReturn(java.util.List.of());
+        ReflectionTestUtils.setField(service, "systemResource", new ByteArrayResource(
+                "skills={dynamically_injected_skills_list}; conversation={current_conversation_id}"
+                        .getBytes(StandardCharsets.UTF_8)));
+        ReflectionTestUtils.setField(service, "toolCallbackProvider", ToolCallbackProvider.from(delegate));
+
+        ChatRequestContext context = ChatRequestContext.builder()
+                .message("question")
+                .conversationId(42L)
+                .subject(subject)
+                .build();
+        service.streamChat(context).collectList().block();
+
+        ArgumentCaptor<Object[]> callbacksCaptor = ArgumentCaptor.forClass(Object[].class);
+        ArgumentCaptor<Map<String, Object>> contextCaptor = ArgumentCaptor.forClass(Map.class);
+        verify(requestSpec).tools(callbacksCaptor.capture());
+        verify(requestSpec).toolContext(contextCaptor.capture());
+        assertInstanceOf(SecurityContextToolCallback.class, callbacksCaptor.getValue()[0]);
+        assertSame(subject, McpContextHolder.getSubject(new org.springframework.ai.chat.model.ToolContext(
+                contextCaptor.getValue())));
+    }
+
     private GeneralConfigDao configDao(AtomicReference<GeneralConfig> currentConfig) {
         return (GeneralConfigDao) Proxy.newProxyInstance(
                 GeneralConfigDao.class.getClassLoader(),

hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java

@@ -23,6 +23,8 @@ import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import com.usthe.sureness.subject.SubjectSum;
+import com.usthe.sureness.util.SurenessContextHolder;
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
@@ -34,6 +36,7 @@ import org.apache.hertzbeat.ai.service.ChatClientProviderService;
 import org.apache.hertzbeat.common.entity.ai.ChatConversation;
 import org.apache.hertzbeat.common.entity.ai.ChatMessage;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.ArgumentCaptor;
 import org.mockito.InjectMocks;
@@ -62,8 +65,15 @@ class ConversationServiceImplTest {
     @InjectMocks
     private ConversationServiceImpl conversationService;
 
+    @AfterEach
+    void clearSecurityContext() {
+        SurenessContextHolder.clear();
+    }
+
     @Test
     void streamChatShouldKeepCompleteConversationHistory() {
+        SubjectSum subject = org.mockito.Mockito.mock(SubjectSum.class);
+        SurenessContextHolder.bindSubject(subject);
         ChatConversation conversation = ChatConversation.builder()
             .id(CONVERSATION_ID)
             .title("已命名会话")
@@ -104,5 +114,6 @@ class ConversationServiceImplTest {
         ArgumentCaptor<ChatRequestContext> contextCaptor = ArgumentCaptor.forClass(ChatRequestContext.class);
         verify(chatClientProviderService).streamChat(contextCaptor.capture());
         assertEquals(history, contextCaptor.getValue().getConversationHistory());
+        assertEquals(subject, contextCaptor.getValue().getSubject());
     }
 }