This is an automated email from the ASF dual-hosted git repository.
yuluo-yx pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/hertzbeat.git
The following commit(s) were added to refs/heads/master by this push:
new 032a4852f7 fix(ai): propagate security context to tool callbacks
(#4230)
032a4852f7 is described below
commit 032a4852f71b7113ba2654ec5444b9b5ea86edec
Author: shown <[email protected]>
AuthorDate: Wed Jul 22 23:46:33 2026 +0800
fix(ai): propagate security context to tool callbacks (#4230)
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]
