maintenance: scope AI conversations by creator (#4280)

Co-authored-by: Duansg <siguoduan@gmail.com>
This commit is contained in:
Logic
2026-08-25 23:05:12 +08:00
committed by GitHub
co-authored by Duansg
parent 7a8156a2d0
commit dc10bcef97
18 changed files with 694 additions and 76 deletions
@@ -17,6 +17,8 @@
package org.apache.hertzbeat.ai.dao; package org.apache.hertzbeat.ai.dao;
import java.util.List;
import java.util.Optional;
import org.apache.hertzbeat.common.entity.ai.ChatConversation; import org.apache.hertzbeat.common.entity.ai.ChatConversation;
import org.springframework.data.jpa.repository.JpaRepository; import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.stereotype.Repository; import org.springframework.stereotype.Repository;
@@ -26,4 +28,8 @@ import org.springframework.stereotype.Repository;
*/ */
@Repository @Repository
public interface ChatConversationDao extends JpaRepository<ChatConversation, Long> { public interface ChatConversationDao extends JpaRepository<ChatConversation, Long> {
Optional<ChatConversation> findByIdAndCreator(Long id, String creator);
List<ChatConversation> findAllByCreatorOrderByIdDesc(String creator);
} }
@@ -19,6 +19,7 @@ package org.apache.hertzbeat.ai.dao;
import java.time.LocalDateTime; import java.time.LocalDateTime;
import java.util.List; import java.util.List;
import java.util.Optional;
import org.apache.hertzbeat.common.entity.ai.SopSchedule; import org.apache.hertzbeat.common.entity.ai.SopSchedule;
import org.springframework.data.jpa.repository.JpaRepository; import org.springframework.data.jpa.repository.JpaRepository;
import org.springframework.data.jpa.repository.JpaSpecificationExecutor; import org.springframework.data.jpa.repository.JpaSpecificationExecutor;
@@ -37,14 +38,23 @@ public interface SopScheduleDao extends JpaRepository<SopSchedule, Long>, JpaSpe
* @param conversationId The conversation ID * @param conversationId The conversation ID
* @return List of schedules * @return List of schedules
*/ */
List<SopSchedule> findByConversationId(Long conversationId); List<SopSchedule> findByConversationIdAndCreator(Long conversationId, String creator);
/**
* Find a schedule only when it belongs to the supplied creator.
* @param id schedule identity
* @param creator authenticated creator
* @return matching schedule
*/
Optional<SopSchedule> findByIdAndCreator(Long id, String creator);
/** /**
* Find all enabled schedules that are due for execution. * Find all enabled schedules that are due for execution.
* @param currentTime The current time to compare against * @param currentTime The current time to compare against
* @return List of due schedules * @return List of due schedules
*/ */
@Query("SELECT s FROM SopSchedule s WHERE s.enabled = true AND s.nextRunTime <= :currentTime") @Query("SELECT s FROM SopSchedule s "
+ "WHERE s.enabled = true AND s.creator IS NOT NULL AND s.nextRunTime <= :currentTime")
List<SopSchedule> findDueSchedules(@Param("currentTime") LocalDateTime currentTime); List<SopSchedule> findDueSchedules(@Param("currentTime") LocalDateTime currentTime);
/** /**
@@ -21,6 +21,7 @@ import java.time.LocalDateTime;
import java.util.HashMap; import java.util.HashMap;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Objects;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import org.apache.hertzbeat.ai.dao.ChatMessageDao; import org.apache.hertzbeat.ai.dao.ChatMessageDao;
import org.apache.hertzbeat.ai.service.SopScheduleService; import org.apache.hertzbeat.ai.service.SopScheduleService;
@@ -98,6 +99,12 @@ public class SopScheduleExecutor {
* Execute a single scheduled SOP and push result to conversation. * Execute a single scheduled SOP and push result to conversation.
*/ */
private void executeSchedule(SopSchedule schedule) { private void executeSchedule(SopSchedule schedule) {
schedule = sopScheduleService.getScheduleForExecution(schedule.getId());
if (schedule == null) {
return;
}
String executionCreator = schedule.getCreator();
Long executionConversationId = schedule.getConversationId();
log.info("Executing scheduled SOP {} for conversation {}", log.info("Executing scheduled SOP {} for conversation {}",
schedule.getSopName(), schedule.getConversationId()); schedule.getSopName(), schedule.getConversationId());
@@ -123,6 +130,12 @@ public class SopScheduleExecutor {
// Execute SOP // Execute SOP
SopResult result = sopEngine.executeSync(definition, params); SopResult result = sopEngine.executeSync(definition, params);
SopSchedule deliverySchedule = sopScheduleService.getScheduleForExecution(schedule.getId());
if (!hasSameExecutionTarget(deliverySchedule, executionCreator, executionConversationId)) {
log.warn("Schedule {} lost its execution owner before result delivery", schedule.getId());
return;
}
// Create push message // Create push message
String messageContent = formatPushMessage(schedule, result); String messageContent = formatPushMessage(schedule, result);
@@ -131,6 +144,7 @@ public class SopScheduleExecutor {
.conversationId(schedule.getConversationId()) .conversationId(schedule.getConversationId())
.role(ROLE_SYSTEM_PUSH) .role(ROLE_SYSTEM_PUSH)
.content(messageContent) .content(messageContent)
.creator(schedule.getCreator())
.build(); .build();
chatMessageDao.save(pushMessage); chatMessageDao.save(pushMessage);
@@ -142,6 +156,12 @@ public class SopScheduleExecutor {
log.error("Failed to execute scheduled SOP {} for conversation {}", log.error("Failed to execute scheduled SOP {} for conversation {}",
schedule.getSopName(), schedule.getConversationId(), e); schedule.getSopName(), schedule.getConversationId(), e);
SopSchedule deliverySchedule = sopScheduleService.getScheduleForExecution(schedule.getId());
if (!hasSameExecutionTarget(deliverySchedule, executionCreator, executionConversationId)) {
log.warn("Schedule {} lost its execution owner before error delivery", schedule.getId());
return;
}
// Still save an error message // Still save an error message
String errorContent = SopMessageUtil.getMessage("schedule.push.error.prefix") + " " + schedule.getSopName() String errorContent = SopMessageUtil.getMessage("schedule.push.error.prefix") + " " + schedule.getSopName()
+ "\n\n" + SopMessageUtil.getMessage("schedule.push.error.label") + " " + e.getMessage(); + "\n\n" + SopMessageUtil.getMessage("schedule.push.error.label") + " " + e.getMessage();
@@ -149,6 +169,7 @@ public class SopScheduleExecutor {
.conversationId(schedule.getConversationId()) .conversationId(schedule.getConversationId())
.role(ROLE_SYSTEM_PUSH) .role(ROLE_SYSTEM_PUSH)
.content(errorContent) .content(errorContent)
.creator(schedule.getCreator())
.build(); .build();
chatMessageDao.save(errorMessage); chatMessageDao.save(errorMessage);
@@ -158,6 +179,12 @@ public class SopScheduleExecutor {
} }
} }
private boolean hasSameExecutionTarget(SopSchedule schedule, String creator, Long conversationId) {
return schedule != null
&& Objects.equals(creator, schedule.getCreator())
&& Objects.equals(conversationId, schedule.getConversationId());
}
/** /**
* Format the push message content with SOP result. * Format the push message content with SOP result.
*/ */
@@ -73,6 +73,15 @@ public interface SopScheduleService {
*/ */
List<SopSchedule> getDueSchedules(); List<SopSchedule> getDueSchedules();
/**
* Re-read a schedule for background execution and verify that its persisted
* creator still owns the target conversation. This method does not depend
* on a request-thread subject.
* @param id schedule ID
* @return validated schedule, or {@code null} when it must not execute
*/
SopSchedule getScheduleForExecution(Long id);
/** /**
* Update the execution times after a schedule runs. * Update the execution times after a schedule runs.
* @param id The schedule ID * @param id The schedule ID
@@ -33,7 +33,6 @@ import org.apache.hertzbeat.common.entity.ai.ChatConversation;
import org.apache.hertzbeat.common.entity.ai.ChatMessage; import org.apache.hertzbeat.common.entity.ai.ChatMessage;
import org.apache.hertzbeat.common.util.AesUtil; import org.apache.hertzbeat.common.util.AesUtil;
import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.domain.Sort;
import org.springframework.http.codec.ServerSentEvent; import org.springframework.http.codec.ServerSentEvent;
import org.springframework.stereotype.Service; import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional; import org.springframework.transaction.annotation.Transactional;
@@ -66,6 +65,10 @@ public class ConversationServiceImpl implements ConversationService {
@Override @Override
public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(String message, Long conversationId) { public Flux<ServerSentEvent<ChatResponseChunk>> streamChat(String message, Long conversationId) {
String creator = requireCurrentUserId();
ChatConversation conversation = conversationId == null
? null
: requireOwnedConversation(conversationId, creator);
// Check if provider is properly configured // Check if provider is properly configured
if (!chatClientProviderService.isConfigured()) { if (!chatClientProviderService.isConfigured()) {
@@ -78,15 +81,12 @@ public class ConversationServiceImpl implements ConversationService {
.build()); .build());
} }
ChatConversation conversation; if (conversation == null) {
if (conversationId == null) {
// The API contract makes conversationId optional, so create a conversation for the first message. // The API contract makes conversationId optional, so create a conversation for the first message.
conversation = new ChatConversation(); conversation = new ChatConversation();
conversation.setTitle(buildConversationTitle(message)); conversation.setTitle(buildConversationTitle(message));
conversation.setCreator(creator);
conversation = conversationDao.save(conversation); conversation = conversationDao.save(conversation);
} else {
conversation = conversationDao.findById(conversationId)
.orElseThrow(() -> new IllegalArgumentException("Conversation not found: " + conversationId));
} }
Long currentConversationId = conversation.getId(); Long currentConversationId = conversation.getId();
log.info("Starting streaming conversation: {}", currentConversationId); log.info("Starting streaming conversation: {}", currentConversationId);
@@ -170,6 +170,7 @@ public class ConversationServiceImpl implements ConversationService {
public ChatConversation createConversation() { public ChatConversation createConversation() {
ChatConversation conversation = new ChatConversation(); ChatConversation conversation = new ChatConversation();
conversation.setTitle("conversation-" + UUID.randomUUID().toString().substring(0, 4)); conversation.setTitle("conversation-" + UUID.randomUUID().toString().substring(0, 4));
conversation.setCreator(requireCurrentUserId());
return conversationDao.save(conversation); return conversationDao.save(conversation);
} }
@@ -182,17 +183,16 @@ public class ConversationServiceImpl implements ConversationService {
if (conversationId == null) { if (conversationId == null) {
return null; return null;
} }
ChatConversation conversation = conversationDao.findById(conversationId).orElse(null); ChatConversation conversation = requireOwnedConversation(conversationId, requireCurrentUserId());
if (conversation != null) { List<ChatMessage> messages = messageDao.findByConversationIdOrderByGmtCreateAsc(conversationId);
List<ChatMessage> messages = messageDao.findByConversationIdOrderByGmtCreateAsc(conversationId); conversation.setMessages(messages);
conversation.setMessages(messages);
}
return conversation; return conversation;
} }
@Override @Override
public List<ChatConversation> getAllConversations() { public List<ChatConversation> getAllConversations() {
List<ChatConversation> conversations = conversationDao.findAll(Sort.by(Sort.Direction.DESC, "id")); List<ChatConversation> conversations =
conversationDao.findAllByCreatorOrderByIdDesc(requireCurrentUserId());
if (conversations.isEmpty()) { if (conversations.isEmpty()) {
return conversations; return conversations;
} }
@@ -213,6 +213,7 @@ public class ConversationServiceImpl implements ConversationService {
@Override @Override
@Transactional(rollbackFor = Exception.class) @Transactional(rollbackFor = Exception.class)
public void deleteConversation(Long conversationId) { public void deleteConversation(Long conversationId) {
requireOwnedConversation(conversationId, requireCurrentUserId());
// Delete associated schedules first to prevent tasks from writing orphaned messages. // Delete associated schedules first to prevent tasks from writing orphaned messages.
sopScheduleDao.deleteByConversationId(conversationId); sopScheduleDao.deleteByConversationId(conversationId);
List<ChatMessage> messages = messageDao.findByConversationIdOrderByGmtCreateAsc(conversationId); List<ChatMessage> messages = messageDao.findByConversationIdOrderByGmtCreateAsc(conversationId);
@@ -224,7 +225,8 @@ public class ConversationServiceImpl implements ConversationService {
@Override @Override
public Boolean saveSecurityData(SecurityData securityData) { public Boolean saveSecurityData(SecurityData securityData) {
Optional<ChatConversation> chatConversation = conversationDao.findById(securityData.getConversationId()); Optional<ChatConversation> chatConversation = conversationDao.findByIdAndCreator(
securityData.getConversationId(), requireCurrentUserId());
if (chatConversation.isPresent()) { if (chatConversation.isPresent()) {
ChatConversation conversation = chatConversation.get(); ChatConversation conversation = chatConversation.get();
conversation.setSecurityData(AesUtil.aesEncode(securityData.getSecurityData())); conversation.setSecurityData(AesUtil.aesEncode(securityData.getSecurityData()));
@@ -234,4 +236,17 @@ public class ConversationServiceImpl implements ConversationService {
return false; return false;
} }
private String requireCurrentUserId() {
SubjectSum subject = SurenessContextHolder.getBindSubject();
if (subject == null || subject.getPrincipal() == null) {
throw new IllegalStateException("No authenticated user");
}
return String.valueOf(subject.getPrincipal());
}
private ChatConversation requireOwnedConversation(Long conversationId, String creator) {
return conversationDao.findByIdAndCreator(conversationId, creator)
.orElseThrow(() -> new IllegalArgumentException("Conversation not found: " + conversationId));
}
} }
@@ -17,11 +17,16 @@
package org.apache.hertzbeat.ai.service.impl; package org.apache.hertzbeat.ai.service.impl;
import com.usthe.sureness.subject.SubjectSum;
import com.usthe.sureness.util.SurenessContextHolder;
import java.time.LocalDateTime; import java.time.LocalDateTime;
import java.util.List; import java.util.List;
import lombok.extern.slf4j.Slf4j; import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.apache.hertzbeat.ai.dao.ChatConversationDao;
import org.apache.hertzbeat.ai.dao.SopScheduleDao; import org.apache.hertzbeat.ai.dao.SopScheduleDao;
import org.apache.hertzbeat.ai.service.SopScheduleService; import org.apache.hertzbeat.ai.service.SopScheduleService;
import org.apache.hertzbeat.common.entity.ai.ChatConversation;
import org.apache.hertzbeat.common.entity.ai.SopSchedule; import org.apache.hertzbeat.common.entity.ai.SopSchedule;
import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.support.CronExpression; import org.springframework.scheduling.support.CronExpression;
@@ -36,23 +41,32 @@ import org.springframework.transaction.annotation.Transactional;
public class SopScheduleServiceImpl implements SopScheduleService { public class SopScheduleServiceImpl implements SopScheduleService {
private final SopScheduleDao sopScheduleDao; private final SopScheduleDao sopScheduleDao;
private final ChatConversationDao conversationDao;
@Autowired @Autowired
public SopScheduleServiceImpl(SopScheduleDao sopScheduleDao) { public SopScheduleServiceImpl(SopScheduleDao sopScheduleDao,
ChatConversationDao conversationDao) {
this.sopScheduleDao = sopScheduleDao; this.sopScheduleDao = sopScheduleDao;
this.conversationDao = conversationDao;
} }
@Override @Override
@Transactional @Transactional
public SopSchedule createSchedule(SopSchedule schedule) { public SopSchedule createSchedule(SopSchedule schedule) {
// Validate cron expression String creator = requireCurrentUserId();
requireOwnedConversation(schedule.getConversationId(), creator);
validateCronExpression(schedule.getCronExpression()); validateCronExpression(schedule.getCronExpression());
// Calculate next run time SopSchedule persisted = SopSchedule.builder()
schedule.setNextRunTime(calculateNextRunTime(schedule.getCronExpression())); .conversationId(schedule.getConversationId())
schedule.setEnabled(schedule.getEnabled() != null ? schedule.getEnabled() : true); .sopName(schedule.getSopName())
.sopParams(schedule.getSopParams())
SopSchedule saved = sopScheduleDao.save(schedule); .cronExpression(schedule.getCronExpression())
.enabled(schedule.getEnabled() != null ? schedule.getEnabled() : true)
.nextRunTime(calculateNextRunTime(schedule.getCronExpression()))
.creator(creator)
.build();
SopSchedule saved = sopScheduleDao.save(persisted);
log.info("Created schedule {} for conversation {} with SOP {}", log.info("Created schedule {} for conversation {} with SOP {}",
saved.getId(), saved.getConversationId(), saved.getSopName()); saved.getId(), saved.getConversationId(), saved.getSopName());
return saved; return saved;
@@ -61,14 +75,13 @@ public class SopScheduleServiceImpl implements SopScheduleService {
@Override @Override
@Transactional @Transactional
public SopSchedule updateSchedule(SopSchedule schedule) { public SopSchedule updateSchedule(SopSchedule schedule) {
SopSchedule existing = sopScheduleDao.findById(schedule.getId()) SopSchedule existing = sopScheduleDao.findByIdAndCreator(
schedule.getId(), requireCurrentUserId())
.orElseThrow(() -> new IllegalArgumentException("Schedule not found: " + schedule.getId())); .orElseThrow(() -> new IllegalArgumentException("Schedule not found: " + schedule.getId()));
// Update fields
existing.setSopName(schedule.getSopName()); existing.setSopName(schedule.getSopName());
existing.setSopParams(schedule.getSopParams()); existing.setSopParams(schedule.getSopParams());
// If cron expression changed, recalculate next run time
if (!existing.getCronExpression().equals(schedule.getCronExpression())) { if (!existing.getCronExpression().equals(schedule.getCronExpression())) {
validateCronExpression(schedule.getCronExpression()); validateCronExpression(schedule.getCronExpression());
existing.setCronExpression(schedule.getCronExpression()); existing.setCronExpression(schedule.getCronExpression());
@@ -85,29 +98,32 @@ public class SopScheduleServiceImpl implements SopScheduleService {
@Override @Override
@Transactional @Transactional
public void deleteSchedule(Long id) { public void deleteSchedule(Long id) {
SopSchedule schedule = sopScheduleDao.findByIdAndCreator(id, requireCurrentUserId())
.orElseThrow(() -> new IllegalArgumentException("Schedule not found: " + id));
log.info("Deleting schedule {}", id); log.info("Deleting schedule {}", id);
sopScheduleDao.deleteById(id); sopScheduleDao.delete(schedule);
} }
@Override @Override
public SopSchedule getSchedule(Long id) { public SopSchedule getSchedule(Long id) {
return sopScheduleDao.findById(id).orElse(null); return sopScheduleDao.findByIdAndCreator(id, requireCurrentUserId()).orElse(null);
} }
@Override @Override
public List<SopSchedule> getSchedulesByConversation(Long conversationId) { public List<SopSchedule> getSchedulesByConversation(Long conversationId) {
return sopScheduleDao.findByConversationId(conversationId); String creator = requireCurrentUserId();
requireOwnedConversation(conversationId, creator);
return sopScheduleDao.findByConversationIdAndCreator(conversationId, creator);
} }
@Override @Override
@Transactional @Transactional
public SopSchedule toggleSchedule(Long id, boolean enabled) { public SopSchedule toggleSchedule(Long id, boolean enabled) {
SopSchedule schedule = sopScheduleDao.findById(id) SopSchedule schedule = sopScheduleDao.findByIdAndCreator(id, requireCurrentUserId())
.orElseThrow(() -> new IllegalArgumentException("Schedule not found: " + id)); .orElseThrow(() -> new IllegalArgumentException("Schedule not found: " + id));
schedule.setEnabled(enabled); schedule.setEnabled(enabled);
// If enabling, recalculate next run time
if (enabled) { if (enabled) {
schedule.setNextRunTime(calculateNextRunTime(schedule.getCronExpression())); schedule.setNextRunTime(calculateNextRunTime(schedule.getCronExpression()));
} }
@@ -123,8 +139,26 @@ public class SopScheduleServiceImpl implements SopScheduleService {
@Override @Override
@Transactional @Transactional
public void updateAfterExecution(Long id) { public SopSchedule getScheduleForExecution(Long id) {
SopSchedule schedule = sopScheduleDao.findById(id).orElse(null); SopSchedule schedule = sopScheduleDao.findById(id).orElse(null);
if (schedule == null || !Boolean.TRUE.equals(schedule.getEnabled())) {
return null;
}
if (StringUtils.isBlank(schedule.getCreator())
|| conversationDao.findByIdAndCreator(
schedule.getConversationId(), schedule.getCreator()).isEmpty()) {
schedule.setEnabled(false);
sopScheduleDao.save(schedule);
log.warn("Disabled schedule {} because its execution owner is missing", id);
return null;
}
return schedule;
}
@Override
@Transactional
public void updateAfterExecution(Long id) {
SopSchedule schedule = getScheduleForExecution(id);
if (schedule == null) { if (schedule == null) {
return; return;
} }
@@ -161,4 +195,18 @@ public class SopScheduleServiceImpl implements SopScheduleService {
throw new IllegalArgumentException("Failed to calculate next run time: " + cronExpression, e); throw new IllegalArgumentException("Failed to calculate next run time: " + cronExpression, e);
} }
} }
private String requireCurrentUserId() {
SubjectSum subject = SurenessContextHolder.getBindSubject();
if (subject == null || subject.getPrincipal() == null) {
throw new IllegalStateException("No authenticated user");
}
return String.valueOf(subject.getPrincipal());
}
private ChatConversation requireOwnedConversation(Long conversationId, String creator) {
return conversationDao.findByIdAndCreator(conversationId, creator)
.orElseThrow(() ->
new IllegalArgumentException("Conversation not found: " + conversationId));
}
} }
@@ -312,8 +312,16 @@ public class MonitorToolsImpl implements MonitorTools {
// Query and add sensitive parameters // Query and add sensitive parameters
if (conversationId != null) { if (conversationId != null) {
Optional<ChatConversation> chatConversation = conversationDao.findById(conversationId); SubjectSum subject = McpContextHolder.getSubject();
if (chatConversation.isPresent() && StringUtils.isNotEmpty(chatConversation.get().getSecurityData())) { if (subject == null || subject.getPrincipal() == null) {
return "Error: Authenticated conversation context is required";
}
Optional<ChatConversation> chatConversation = conversationDao.findByIdAndCreator(
conversationId, String.valueOf(subject.getPrincipal()));
if (chatConversation.isEmpty()) {
return "Error: Conversation not found or inaccessible";
}
if (StringUtils.isNotEmpty(chatConversation.get().getSecurityData())) {
List<Param> securityParams = JsonUtil.fromJson( List<Param> securityParams = JsonUtil.fromJson(
AesUtil.aesDecode(chatConversation.get().getSecurityData()), AesUtil.aesDecode(chatConversation.get().getSecurityData()),
new TypeReference<List<Param>>() { new TypeReference<List<Param>>() {
@@ -20,6 +20,7 @@ package org.apache.hertzbeat.ai.schedule;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyMap; import static org.mockito.ArgumentMatchers.anyMap;
import static org.mockito.ArgumentMatchers.argThat;
import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.times; import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
@@ -77,6 +78,8 @@ class SopScheduleExecutorTest {
.content("ok") .content("ok")
.build(); .build();
when(scheduleService.getDueSchedules()).thenReturn(List.of(first, second)); when(scheduleService.getDueSchedules()).thenReturn(List.of(first, second));
when(scheduleService.getScheduleForExecution(1L)).thenReturn(first, first);
when(scheduleService.getScheduleForExecution(2L)).thenReturn(second, second);
when(skillRegistry.getSkill("daily_inspection")).thenReturn(definition); when(skillRegistry.getSkill("daily_inspection")).thenReturn(definition);
when(sopEngine.executeSync(any(SopDefinition.class), anyMap())).thenReturn(result); when(sopEngine.executeSync(any(SopDefinition.class), anyMap())).thenReturn(result);
doThrow(new IllegalStateException("database unavailable")) doThrow(new IllegalStateException("database unavailable"))
@@ -85,6 +88,8 @@ class SopScheduleExecutorTest {
executor.checkAndExecuteDueSchedules(); executor.checkAndExecuteDueSchedules();
verify(sopEngine, times(2)).executeSync(any(SopDefinition.class), anyMap()); verify(sopEngine, times(2)).executeSync(any(SopDefinition.class), anyMap());
verify(chatMessageDao, times(2)).save(argThat(
message -> "alice".equals(message.getCreator())));
verify(scheduleService).updateAfterExecution(2L); verify(scheduleService).updateAfterExecution(2L);
} }
@@ -92,6 +97,7 @@ class SopScheduleExecutorTest {
void checkShouldRejectInvalidScheduleParameters() { void checkShouldRejectInvalidScheduleParameters() {
SopSchedule schedule = schedule(1L, "not-json"); SopSchedule schedule = schedule(1L, "not-json");
when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule)); when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule));
when(scheduleService.getScheduleForExecution(1L)).thenReturn(schedule, schedule);
when(skillRegistry.getSkill("daily_inspection")) when(skillRegistry.getSkill("daily_inspection"))
.thenReturn(SopDefinition.builder().name("daily_inspection").build()); .thenReturn(SopDefinition.builder().name("daily_inspection").build());
@@ -102,10 +108,44 @@ class SopScheduleExecutorTest {
verify(scheduleService).updateAfterExecution(1L); verify(scheduleService).updateAfterExecution(1L);
} }
@Test
void checkShouldSkipScheduleWithoutValidatedOwner() {
SopSchedule schedule = schedule(1L, null);
when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule));
when(scheduleService.getScheduleForExecution(1L)).thenReturn(null);
executor.checkAndExecuteDueSchedules();
verifyNoInteractions(sopEngine, chatMessageDao);
verify(scheduleService, times(0)).updateAfterExecution(1L);
}
@Test
void checkShouldNotDeliverWhenOwnerChangesDuringExecution() {
SopSchedule schedule = schedule(1L, null);
SopSchedule changedOwner = schedule(1L, null);
changedOwner.setCreator("bob");
changedOwner.setConversationId(20L);
when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule));
when(scheduleService.getScheduleForExecution(1L))
.thenReturn(schedule, changedOwner, changedOwner);
when(skillRegistry.getSkill("daily_inspection"))
.thenReturn(SopDefinition.builder().name("daily_inspection").build());
when(sopEngine.executeSync(any(SopDefinition.class), anyMap()))
.thenReturn(SopResult.builder().status("SUCCESS").content("ok").build());
executor.checkAndExecuteDueSchedules();
verify(sopEngine).executeSync(any(SopDefinition.class), anyMap());
verifyNoInteractions(chatMessageDao);
verify(scheduleService).updateAfterExecution(1L);
}
@Test @Test
void checkShouldPushErrorWhenScheduledSkillNoLongerExists() { void checkShouldPushErrorWhenScheduledSkillNoLongerExists() {
SopSchedule schedule = schedule(1L, null); SopSchedule schedule = schedule(1L, null);
when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule)); when(scheduleService.getDueSchedules()).thenReturn(List.of(schedule));
when(scheduleService.getScheduleForExecution(1L)).thenReturn(schedule, schedule);
when(skillRegistry.getSkill("daily_inspection")).thenReturn(null); when(skillRegistry.getSkill("daily_inspection")).thenReturn(null);
executor.checkAndExecuteDueSchedules(); executor.checkAndExecuteDueSchedules();
@@ -123,6 +163,7 @@ class SopScheduleExecutorTest {
.conversationId(10L) .conversationId(10L)
.sopName("daily_inspection") .sopName("daily_inspection")
.sopParams(params) .sopParams(params)
.creator("alice")
.build(); .build();
} }
} }
@@ -18,9 +18,13 @@
package org.apache.hertzbeat.ai.service.impl; package org.apache.hertzbeat.ai.service.impl;
import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull; import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when; import static org.mockito.Mockito.when;
@@ -35,6 +39,7 @@ import org.apache.hertzbeat.ai.dao.ChatMessageDao;
import org.apache.hertzbeat.ai.dao.SopScheduleDao; import org.apache.hertzbeat.ai.dao.SopScheduleDao;
import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext; import org.apache.hertzbeat.ai.pojo.dto.ChatRequestContext;
import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk; import org.apache.hertzbeat.ai.pojo.dto.ChatResponseChunk;
import org.apache.hertzbeat.ai.pojo.dto.SecurityData;
import org.apache.hertzbeat.ai.service.ChatClientProviderService; import org.apache.hertzbeat.ai.service.ChatClientProviderService;
import org.apache.hertzbeat.common.entity.ai.ChatConversation; import org.apache.hertzbeat.common.entity.ai.ChatConversation;
import org.apache.hertzbeat.common.entity.ai.ChatMessage; import org.apache.hertzbeat.common.entity.ai.ChatMessage;
@@ -79,29 +84,30 @@ class ConversationServiceImplTest {
@Test @Test
void streamChatShouldKeepCompleteConversationHistory() { void streamChatShouldKeepCompleteConversationHistory() {
SubjectSum subject = org.mockito.Mockito.mock(SubjectSum.class); SubjectSum subject = bindSubject("alice");
SurenessContextHolder.bindSubject(subject);
ChatConversation conversation = ChatConversation.builder() ChatConversation conversation = ChatConversation.builder()
.id(CONVERSATION_ID) .id(CONVERSATION_ID)
.title("已命名会话") .title("Named conversation")
.creator("alice")
.build(); .build();
List<ChatMessage> history = List.of( List<ChatMessage> history = List.of(
ChatMessage.builder() ChatMessage.builder()
.id(11L) .id(11L)
.conversationId(CONVERSATION_ID) .conversationId(CONVERSATION_ID)
.role("user") .role("user")
.content("上一轮问题") .content("Previous question")
.build(), .build(),
ChatMessage.builder() ChatMessage.builder()
.id(12L) .id(12L)
.conversationId(CONVERSATION_ID) .conversationId(CONVERSATION_ID)
.role("assistant") .role("assistant")
.content("上一轮回答") .content("Previous answer")
.build()); .build());
AtomicLong messageId = new AtomicLong(20L); AtomicLong messageId = new AtomicLong(20L);
when(chatClientProviderService.isConfigured()).thenReturn(true); when(chatClientProviderService.isConfigured()).thenReturn(true);
when(conversationDao.findById(CONVERSATION_ID)).thenReturn(Optional.of(conversation)); when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.of(conversation));
when(messageDao.findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID)).thenReturn(history); when(messageDao.findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID)).thenReturn(history);
when(messageDao.save(any(ChatMessage.class))).thenAnswer(invocation -> { when(messageDao.save(any(ChatMessage.class))).thenAnswer(invocation -> {
ChatMessage savedMessage = invocation.getArgument(0); ChatMessage savedMessage = invocation.getArgument(0);
@@ -109,10 +115,10 @@ class ConversationServiceImplTest {
return savedMessage; return savedMessage;
}); });
when(chatClientProviderService.streamChat(any(ChatRequestContext.class))) when(chatClientProviderService.streamChat(any(ChatRequestContext.class)))
.thenReturn(Flux.just("本轮回答")); .thenReturn(Flux.just("Current answer"));
List<ServerSentEvent<ChatResponseChunk>> events = conversationService List<ServerSentEvent<ChatResponseChunk>> events = conversationService
.streamChat("本轮问题", CONVERSATION_ID) .streamChat("Current question", CONVERSATION_ID)
.collectList() .collectList()
.block(); .block();
@@ -129,6 +135,7 @@ class ConversationServiceImplTest {
*/ */
@Test @Test
void streamChatShouldCreateConversationWhenConversationIdIsMissing() { void streamChatShouldCreateConversationWhenConversationIdIsMissing() {
bindSubject("alice");
AtomicLong messageId = new AtomicLong(20L); AtomicLong messageId = new AtomicLong(20L);
when(chatClientProviderService.isConfigured()).thenReturn(true); when(chatClientProviderService.isConfigured()).thenReturn(true);
when(conversationDao.save(any(ChatConversation.class))).thenAnswer(invocation -> { when(conversationDao.save(any(ChatConversation.class))).thenAnswer(invocation -> {
@@ -143,10 +150,10 @@ class ConversationServiceImplTest {
return savedMessage; return savedMessage;
}); });
when(chatClientProviderService.streamChat(any(ChatRequestContext.class))) when(chatClientProviderService.streamChat(any(ChatRequestContext.class)))
.thenReturn(Flux.just("本轮回答")); .thenReturn(Flux.just("Current answer"));
List<ServerSentEvent<ChatResponseChunk>> events = conversationService List<ServerSentEvent<ChatResponseChunk>> events = conversationService
.streamChat("本轮问题", null) .streamChat("Initial question", null)
.collectList() .collectList()
.block(); .block();
@@ -161,7 +168,8 @@ class ConversationServiceImplTest {
assertEquals(List.of(), contextCaptor.getValue().getConversationHistory()); assertEquals(List.of(), contextCaptor.getValue().getConversationHistory());
ArgumentCaptor<ChatConversation> conversationCaptor = ArgumentCaptor.forClass(ChatConversation.class); ArgumentCaptor<ChatConversation> conversationCaptor = ArgumentCaptor.forClass(ChatConversation.class);
verify(conversationDao).save(conversationCaptor.capture()); verify(conversationDao).save(conversationCaptor.capture());
assertEquals("本轮问题", conversationCaptor.getValue().getTitle()); assertEquals("Initial question", conversationCaptor.getValue().getTitle());
assertEquals("alice", conversationCaptor.getValue().getCreator());
verifyNoMoreInteractions(conversationDao); verifyNoMoreInteractions(conversationDao);
} }
@@ -170,12 +178,20 @@ class ConversationServiceImplTest {
*/ */
@Test @Test
void deleteConversationShouldRemoveSchedulesMessagesAndConversationInOrder() { void deleteConversationShouldRemoveSchedulesMessagesAndConversationInOrder() {
bindSubject("alice");
ChatConversation conversation = ChatConversation.builder()
.id(CONVERSATION_ID)
.title("Owned conversation")
.creator("alice")
.build();
ChatMessage message = ChatMessage.builder() ChatMessage message = ChatMessage.builder()
.id(11L) .id(11L)
.conversationId(CONVERSATION_ID) .conversationId(CONVERSATION_ID)
.role("user") .role("user")
.content("message to delete") .content("message to delete")
.build(); .build();
when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.of(conversation));
when(messageDao.findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID)) when(messageDao.findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID))
.thenReturn(List.of(message)); .thenReturn(List.of(message));
@@ -186,4 +202,88 @@ class ConversationServiceImplTest {
deletionOrder.verify(messageDao).deleteAll(List.of(message)); deletionOrder.verify(messageDao).deleteAll(List.of(message));
deletionOrder.verify(conversationDao).deleteById(CONVERSATION_ID); deletionOrder.verify(conversationDao).deleteById(CONVERSATION_ID);
} }
@Test
void listConversationsShouldExcludeOtherCreators() {
bindSubject("alice");
ChatConversation ownedConversation = ChatConversation.builder()
.id(CONVERSATION_ID)
.title("Owned conversation")
.creator("alice")
.build();
when(conversationDao.findAllByCreatorOrderByIdDesc("alice"))
.thenReturn(List.of(ownedConversation));
when(messageDao.findByConversationIdInOrderByGmtCreateAsc(List.of(CONVERSATION_ID)))
.thenReturn(List.of());
List<ChatConversation> result = conversationService.getAllConversations();
assertEquals(List.of(ownedConversation), result);
verify(conversationDao).findAllByCreatorOrderByIdDesc("alice");
}
@Test
void getConversationShouldRejectAnotherCreator() {
bindSubject("alice");
when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class,
() -> conversationService.getConversation(CONVERSATION_ID));
verify(messageDao, never()).findByConversationIdOrderByGmtCreateAsc(CONVERSATION_ID);
}
@Test
void deleteConversationShouldRejectAnotherCreator() {
bindSubject("alice");
when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class,
() -> conversationService.deleteConversation(CONVERSATION_ID));
verify(sopScheduleDao, never()).deleteByConversationId(CONVERSATION_ID);
verify(conversationDao, never()).deleteById(CONVERSATION_ID);
}
@Test
void streamChatShouldRejectAnotherCreator() {
bindSubject("alice");
when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class,
() -> conversationService.streamChat("Current question", CONVERSATION_ID));
verify(messageDao, never()).save(any(ChatMessage.class));
}
@Test
void createConversationShouldRecordCurrentCreator() {
bindSubject("alice");
when(conversationDao.save(any(ChatConversation.class)))
.thenAnswer(invocation -> invocation.getArgument(0));
ChatConversation conversation = conversationService.createConversation();
assertEquals("alice", conversation.getCreator());
}
@Test
void saveSecurityDataShouldRejectAnotherCreator() {
bindSubject("alice");
SecurityData securityData = new SecurityData();
securityData.setConversationId(CONVERSATION_ID);
securityData.setSecurityData("sensitive-value");
when(conversationDao.findByIdAndCreator(CONVERSATION_ID, "alice"))
.thenReturn(Optional.empty());
assertFalse(conversationService.saveSecurityData(securityData));
verify(conversationDao, never()).save(any(ChatConversation.class));
}
private SubjectSum bindSubject(String principal) {
SubjectSum subject = mock(SubjectSum.class);
when(subject.getPrincipal()).thenReturn(principal);
SurenessContextHolder.bindSubject(subject);
return subject;
}
} }
@@ -1,58 +1,186 @@
/* /*
* Licensed to the Apache Software Foundation (ASF) under one or more * Licensed to the Apache Software Foundation (ASF) under one
* contributor license agreements. See the NOTICE file distributed with * or more contributor license agreements. See the NOTICE file
* this work for additional information regarding copyright ownership. * distributed with this work for additional information
* The ASF licenses this file to You under the Apache License, Version 2.0 * regarding copyright ownership. The ASF licenses this file
* (the "License"); you may not use this file except in compliance with * to You under the Apache License, Version 2.0 (the
* the License. You may obtain a copy of the License at * "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 * http://www.apache.org/licenses/LICENSE-2.0
* *
* Unless required by applicable law or agreed to in writing, software * Unless required by applicable law or agreed to in writing,
* distributed under the License is distributed on an "AS IS" BASIS, * software distributed under the License is distributed on an
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* See the License for the specific language governing permissions and * KIND, either express or implied. See the License for the
* limitations under the License. * specific language governing permissions and limitations
* under the License.
*/ */
package org.apache.hertzbeat.ai.service.impl; package org.apache.hertzbeat.ai.service.impl;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
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.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue; import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.lenient;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions; import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
import com.usthe.sureness.subject.SubjectSum;
import com.usthe.sureness.util.SurenessContextHolder;
import java.util.Optional;
import org.apache.hertzbeat.ai.dao.ChatConversationDao;
import org.apache.hertzbeat.ai.dao.SopScheduleDao; import org.apache.hertzbeat.ai.dao.SopScheduleDao;
import org.apache.hertzbeat.common.entity.ai.ChatConversation;
import org.apache.hertzbeat.common.entity.ai.SopSchedule; import org.apache.hertzbeat.common.entity.ai.SopSchedule;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test; import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith; import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock; import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoExtension;
/** /**
* Verifies that SOP schedules with no future execution time are not persisted. * Ownership and scheduling contracts for user-facing SOP schedule operations.
*/ */
@ExtendWith(MockitoExtension.class) @ExtendWith(MockitoExtension.class)
class SopScheduleServiceImplTest { class SopScheduleServiceImplTest {
@Mock @Mock
private SopScheduleDao sopScheduleDao; private SopScheduleDao scheduleDao;
@InjectMocks @Mock
private SopScheduleServiceImpl scheduleService; private ChatConversationDao conversationDao;
private SopScheduleServiceImpl service;
@BeforeEach
void setUp() {
service = new SopScheduleServiceImpl(scheduleDao, conversationDao);
SubjectSum subject = mock(SubjectSum.class);
lenient().when(subject.getPrincipal()).thenReturn("alice");
SurenessContextHolder.bindSubject(subject);
}
@AfterEach
void clearSubject() {
SurenessContextHolder.clear();
}
@Test
void createShouldNotTrustRequestCreator() {
SopSchedule request = schedule(1L, "bob");
when(conversationDao.findByIdAndCreator(10L, "alice"))
.thenReturn(Optional.of(ChatConversation.builder()
.id(10L)
.creator("alice")
.build()));
when(scheduleDao.save(any(SopSchedule.class)))
.thenAnswer(invocation -> invocation.getArgument(0));
SopSchedule created = service.createSchedule(request);
assertEquals("alice", created.getCreator());
}
@Test
void getShouldHideAnotherCreatorsSchedule() {
when(scheduleDao.findByIdAndCreator(1L, "alice")).thenReturn(Optional.empty());
assertNull(service.getSchedule(1L));
}
@Test
void listShouldRejectAnotherCreatorsConversation() {
when(conversationDao.findByIdAndCreator(10L, "alice")).thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class,
() -> service.getSchedulesByConversation(10L));
verify(scheduleDao, never()).findByConversationIdAndCreator(10L, "alice");
}
@Test
void deleteShouldNotRemoveAnotherCreatorsSchedule() {
when(scheduleDao.findByIdAndCreator(1L, "alice")).thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class, () -> service.deleteSchedule(1L));
verify(scheduleDao, never()).delete(any(SopSchedule.class));
}
@Test
void updateAndToggleShouldNotModifyAnotherCreatorsSchedule() {
when(scheduleDao.findByIdAndCreator(1L, "alice")).thenReturn(Optional.empty());
assertThrows(IllegalArgumentException.class,
() -> service.updateSchedule(schedule(1L, "bob")));
assertThrows(IllegalArgumentException.class,
() -> service.toggleSchedule(1L, true));
verify(scheduleDao, never()).save(any(SopSchedule.class));
}
@Test
void backgroundExecutionShouldDisableMissingOwner() {
SopSchedule schedule = schedule(1L, "legacy-owner");
schedule.setEnabled(true);
when(scheduleDao.findById(1L)).thenReturn(Optional.of(schedule));
when(conversationDao.findByIdAndCreator(10L, "legacy-owner"))
.thenReturn(Optional.empty());
when(scheduleDao.save(schedule)).thenReturn(schedule);
assertNull(service.getScheduleForExecution(1L));
assertFalse(schedule.getEnabled());
verify(scheduleDao).save(schedule);
}
@Test
void backgroundExecutionUsesPersistedOwnerWithoutRequestSubject() {
SurenessContextHolder.clear();
SopSchedule schedule = schedule(1L, "alice");
schedule.setEnabled(true);
when(scheduleDao.findById(1L)).thenReturn(Optional.of(schedule));
when(conversationDao.findByIdAndCreator(10L, "alice"))
.thenReturn(Optional.of(ChatConversation.builder()
.id(10L)
.creator("alice")
.build()));
assertSame(schedule, service.getScheduleForExecution(1L));
}
@Test @Test
void createScheduleShouldRejectCronWithoutFutureExecutionTime() { void createScheduleShouldRejectCronWithoutFutureExecutionTime() {
SopSchedule schedule = SopSchedule.builder() SopSchedule schedule = SopSchedule.builder()
.conversationId(1L) .conversationId(10L)
.sopName("daily_inspection") .sopName("daily_inspection")
.cronExpression("0 0 0 31 2 *") .cronExpression("0 0 0 31 2 *")
.build(); .build();
when(conversationDao.findByIdAndCreator(10L, "alice"))
.thenReturn(Optional.of(ChatConversation.builder()
.id(10L)
.creator("alice")
.build()));
IllegalArgumentException exception = assertThrows( IllegalArgumentException exception = assertThrows(
IllegalArgumentException.class, () -> scheduleService.createSchedule(schedule)); IllegalArgumentException.class, () -> service.createSchedule(schedule));
assertTrue(exception.getMessage().contains("no future execution time")); assertTrue(exception.getMessage().contains("no future execution time"));
verifyNoInteractions(sopScheduleDao); verifyNoInteractions(scheduleDao);
}
private SopSchedule schedule(Long id, String creator) {
return SopSchedule.builder()
.id(id)
.conversationId(10L)
.sopName("daily_inspection")
.cronExpression("0 0 9 * * ?")
.creator(creator)
.build();
} }
} }
@@ -0,0 +1,77 @@
/*
* 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.tools.impl;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
import static org.mockito.Mockito.when;
import com.usthe.sureness.subject.SubjectSum;
import java.util.Optional;
import org.apache.hertzbeat.ai.config.McpContextHolder;
import org.apache.hertzbeat.ai.dao.ChatConversationDao;
import org.apache.hertzbeat.manager.service.AppService;
import org.apache.hertzbeat.manager.service.MonitorService;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.InjectMocks;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
/**
* Verifies that protected monitor creation cannot load another user's
* conversation credentials.
*/
@ExtendWith(MockitoExtension.class)
class MonitorToolsImplTest {
@Mock
private MonitorService monitorService;
@Mock
private AppService appService;
@Mock
private ChatConversationDao conversationDao;
@InjectMocks
private MonitorToolsImpl monitorTools;
@AfterEach
void clearContext() {
McpContextHolder.clear();
}
@Test
void protectedAddShouldRejectConversationOutsideCurrentCreator() {
SubjectSum subject = mock(SubjectSum.class);
when(subject.getPrincipal()).thenReturn("alice");
McpContextHolder.setSubject(subject);
when(conversationDao.findByIdAndCreator(10L, "alice")).thenReturn(Optional.empty());
String result = monitorTools.addMonitorProtected(
10L, "database", "mysql", 60, "{\"host\":\"db.local\"}", null);
assertEquals("Error: Conversation not found or inaccessible", result);
verify(conversationDao).findByIdAndCreator(10L, "alice");
verifyNoInteractions(monitorService);
}
}
@@ -19,12 +19,14 @@ package org.apache.hertzbeat.common.entity.ai;
import static io.swagger.v3.oas.annotations.media.Schema.AccessMode.READ_ONLY; import static io.swagger.v3.oas.annotations.media.Schema.AccessMode.READ_ONLY;
import com.fasterxml.jackson.annotation.JsonIgnore;
import io.swagger.v3.oas.annotations.media.Schema; import io.swagger.v3.oas.annotations.media.Schema;
import jakarta.persistence.Entity; import jakarta.persistence.Entity;
import jakarta.persistence.EntityListeners; import jakarta.persistence.EntityListeners;
import jakarta.persistence.GeneratedValue; import jakarta.persistence.GeneratedValue;
import jakarta.persistence.GenerationType; import jakarta.persistence.GenerationType;
import jakarta.persistence.Id; import jakarta.persistence.Id;
import jakarta.persistence.Index;
import jakarta.persistence.OneToMany; import jakarta.persistence.OneToMany;
import jakarta.persistence.Table; import jakarta.persistence.Table;
import org.springframework.data.annotation.CreatedBy; import org.springframework.data.annotation.CreatedBy;
@@ -47,7 +49,9 @@ import java.util.List;
@Builder @Builder
@Entity @Entity
@EntityListeners(AuditingEntityListener.class) @EntityListeners(AuditingEntityListener.class)
@Table(name = "hzb_ai_conversation") @Table(name = "hzb_ai_conversation", indexes = {
@Index(name = "idx_ai_conversation_creator", columnList = "creator")
})
@AllArgsConstructor @AllArgsConstructor
@NoArgsConstructor @NoArgsConstructor
public class ChatConversation { public class ChatConversation {
@@ -81,5 +85,6 @@ public class ChatConversation {
@OneToMany(mappedBy = "conversation") @OneToMany(mappedBy = "conversation")
private List<ChatMessage> messages; private List<ChatMessage> messages;
@JsonIgnore
private String securityData; private String securityData;
} }
@@ -51,6 +51,7 @@ import org.springframework.data.jpa.domain.support.AuditingEntityListener;
@EntityListeners(AuditingEntityListener.class) @EntityListeners(AuditingEntityListener.class)
@Table(name = "hzb_sop_schedule", indexes = { @Table(name = "hzb_sop_schedule", indexes = {
@Index(name = "idx_schedule_conversation_id", columnList = "conversation_id"), @Index(name = "idx_schedule_conversation_id", columnList = "conversation_id"),
@Index(name = "idx_schedule_creator_conversation", columnList = "creator, conversation_id"),
@Index(name = "idx_schedule_enabled_next", columnList = "enabled, next_run_time") @Index(name = "idx_schedule_enabled_next", columnList = "enabled, next_run_time")
}) })
@AllArgsConstructor @AllArgsConstructor
@@ -0,0 +1,45 @@
/*
* 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.common.entity.ai;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import org.apache.hertzbeat.common.util.JsonUtil;
import org.junit.jupiter.api.Test;
/**
* Tests AI conversation serialization.
*/
class ChatConversationTest {
@Test
void serializationShouldExcludeStoredSecurityData() {
ChatConversation conversation = ChatConversation.builder()
.id(1L)
.title("Owned conversation")
.securityData("encrypted-value")
.build();
String json = JsonUtil.toJson(conversation);
assertNotNull(json);
assertFalse(json.contains("securityData"));
assertFalse(json.contains("encrypted-value"));
}
}
@@ -0,0 +1,24 @@
-- 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.
UPDATE hzb_sop_schedule
SET enabled = 0
WHERE creator IS NULL
OR TRIM(creator) = '';
CREATE INDEX IF NOT EXISTS idx_schedule_creator_conversation
ON hzb_sop_schedule(creator, conversation_id);
@@ -0,0 +1,24 @@
-- 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.
UPDATE hzb_sop_schedule
SET enabled = 0
WHERE creator IS NULL
OR TRIM(creator) = '';
CREATE INDEX idx_schedule_creator_conversation
ON hzb_sop_schedule(creator, conversation_id);
@@ -0,0 +1,24 @@
-- 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.
UPDATE hzb_sop_schedule
SET enabled = 0
WHERE creator IS NULL
OR BTRIM(creator) = '';
CREATE INDEX idx_schedule_creator_conversation
ON hzb_sop_schedule(creator, conversation_id);
+26
View File
@@ -129,4 +129,30 @@ When `warehouse.store.greptime.enabled=true`, HertzBeat writes two different kin
- `bin/shutdown.sh` stops the HertzBeat process and downloads the new installation package - `bin/shutdown.sh` stops the HertzBeat process and downloads the new installation package
- Refer to [Installation package to install HertzBeat](./package-deploy) to start with the new installation package and configure the database connection in `application.yml` - Refer to [Installation package to install HertzBeat](./package-deploy) to start with the new installation package and configure the database connection in `application.yml`
## AI Schedule Ownership After Upgrade
AI conversations without a recorded creator are isolated and do not appear in
any user's conversation list. Scheduled AI SOP tasks are owned by the creator
of their target conversation. During upgrade, schedules without a recorded
creator are disabled. Schedules without a target conversation or whose creator
does not match the conversation creator are disabled before they can execute.
The records remain in the database so an administrator can recover them after
verifying the intended owner.
Ownerless schedules are disabled by the database migration. Missing
conversations and creator mismatches are rechecked and disabled before every
background execution.
Before re-enabling a legacy schedule:
1. Back up the metadata database.
2. Verify the owner of the target row in `hzb_ai_conversation`.
3. Set the same verified principal in the conversation and schedule `creator`
columns.
4. Re-enable only the reviewed schedule.
Do not assign all legacy rows to a shared account. A schedule is executed only
while its stored creator still owns the target conversation; ownership
mismatches are disabled automatically.
**HAVE FUN** **HAVE FUN**