This is an automated email from the ASF dual-hosted git repository. yuluo-yx pushed a commit to branch 0722-yuluo/fix-ai-tool-security-context in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
commit 1c9b25a9314aa9dbeac0e81ee25357157c4404e0 Author: yuluo-yx <[email protected]> AuthorDate: Wed Jul 22 22:51:02 2026 +0800 fix(ai): propagate security context to tool callbacks Signed-off-by: yuluo-yx <[email protected]> --- .../hertzbeat/ai/config/McpContextHolder.java | 74 ++++++++++- .../ai/config/SecurityContextToolCallback.java | 60 +++++++++ .../hertzbeat/ai/controller/ChatController.java | 5 - .../hertzbeat/ai/pojo/dto/ChatRequestContext.java | 12 +- .../impl/ChatClientProviderServiceImpl.java | 10 +- .../ai/service/impl/ConversationServiceImpl.java | 4 + .../hertzbeat/ai/config/McpContextHolderTest.java | 146 +++++++++++++++++++++ .../impl/ChatClientProviderServiceImplTest.java | 67 ++++++++++ .../service/impl/ConversationServiceImplTest.java | 11 ++ 9 files changed, 379 insertions(+), 10 deletions(-) diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/McpContextHolder.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/McpContextHolder.java index a41b8f3b23..6e47aa5042 100644 --- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/McpContextHolder.java +++ b/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); + } + } } diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/SecurityContextToolCallback.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/config/SecurityContextToolCallback.java new file mode 100644 index 0000000000..3d193ae7a2 --- /dev/null +++ b/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)); + } +} diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/controller/ChatController.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/controller/ChatController.java index 0f5bfcd9e2..fe8d566096 100644 --- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/controller/ChatController.java +++ b/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()) diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/pojo/dto/ChatRequestContext.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/pojo/dto/ChatRequestContext.java index d0dcfacf06..b7db5be2de 100644 --- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/pojo/dto/ChatRequestContext.java +++ b/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; } diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImpl.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImpl.java index 220a95c816..860b70f6b8 100644 --- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImpl.java +++ b/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())) diff --git a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java b/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java index b642d62e5e..1adbe6f78c 100644 --- a/hertzbeat-ai/src/main/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImpl.java +++ b/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); diff --git a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/config/McpContextHolderTest.java b/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/config/McpContextHolderTest.java new file mode 100644 index 0000000000..2f8b343f72 --- /dev/null +++ b/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()); + } +} diff --git a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImplTest.java b/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImplTest.java index d1f5f2b4bc..c98cce3d46 100644 --- a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ChatClientProviderServiceImplTest.java +++ b/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(), diff --git a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java b/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java index f46c5dcdc9..ce5f5bf3ef 100644 --- a/hertzbeat-ai/src/test/java/org/apache/hertzbeat/ai/service/impl/ConversationServiceImplTest.java +++ b/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()); } } --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
