diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandler.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandler.java index 0d11b4cbea..ab1eea1dea 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandler.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandler.java @@ -150,7 +150,7 @@ public class AgentAlertAnalysisEventHandler { .build(); InvokeCommand command = InvokeCommand.builder() .envelope(GatewayEnvelope.builder() - .channelId(ChannelId.ALERT.id()) + .channelId(ChannelId.SYSTEM.id()) .receivedAt(now) .actor(AgentActor.alertAnalysisActor()) .preferredLanguage(AgentResponseLanguage.systemDefault()) diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandService.java index 6569551f40..7ccaadcf47 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandService.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandService.java @@ -39,6 +39,7 @@ import org.apache.hertzbeat.ai.gateway.conversation.AgentTranscriptRecorder; import org.apache.hertzbeat.ai.gateway.runtime.AgentApprovalHandling; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEvent; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEventType; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeRequest; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeService; import org.apache.hertzbeat.ai.gateway.runtime.TranscriptMessage; @@ -118,13 +119,14 @@ public class AgentCommandService { AgentRuntimeRequest prepare(GatewayCommand command, UserInput userInput) { GatewayEnvelope envelope = command.envelope(); - AgentSession session = sessionService.findOrCreateSession(envelope, userInput); - AgentRun run = runService.createOrResumeRun(session, userInput); + AgentRuntimeEntryType entryType = ((InvokeCommand) command).entryType(); + AgentSession session = sessionService.findOrCreateSession(envelope, userInput, entryType); + AgentRun run = runService.createOrResumeRun(session, userInput, entryType); List chatHistory = transcriptRecorder.chatHistory(session.getId()); transcriptRecorder.recordUserTranscriptEntry(session, run, userInput); AgentRun runningRun = runService.markRunning(run); return AgentRuntimeRequest.builder() - .entryType(((InvokeCommand) command).entryType()) + .entryType(entryType) .approvalHandling(command.replyMode() == ReplyMode.STREAM ? AgentApprovalHandling.WAIT_FOR_DECISION : AgentApprovalHandling.DENY) diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommand.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommand.java index 2feaf7577a..ba609bb344 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommand.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommand.java @@ -51,7 +51,7 @@ public sealed interface GatewayCommand permits String commandId(); /** - * Normalized runtime invocation shared by interactive and alert channels. + * Normalized runtime invocation shared by interactive and system-triggered entries. */ @Builder record InvokeCommand( @@ -125,11 +125,14 @@ public sealed interface GatewayCommand permits GatewayEnvelope envelope, ReplyMode replyMode, String commandId, + AgentRuntimeEntryType originEntryType, String sessionUid) implements GatewayCommand { public GetSessionCommand { envelope = Objects.requireNonNull(envelope, "envelope is required"); replyMode = Objects.requireNonNull(replyMode, "replyMode is required"); + // The command must fully describe the session query scope. + originEntryType = Objects.requireNonNull(originEntryType, "originEntryType is required"); if (!StringUtils.hasText(commandId) || !StringUtils.hasText(sessionUid)) { throw new IllegalArgumentException("commandId and sessionUid are required"); } @@ -147,6 +150,7 @@ public sealed interface GatewayCommand permits GatewayEnvelope envelope, ReplyMode replyMode, String commandId, + AgentRuntimeEntryType originEntryType, String title, int pageIndex, int pageSize) implements GatewayCommand { @@ -154,6 +158,8 @@ public sealed interface GatewayCommand permits public ListSessionsCommand { envelope = Objects.requireNonNull(envelope, "envelope is required"); replyMode = Objects.requireNonNull(replyMode, "replyMode is required"); + // The command must fully describe the session query scope. + originEntryType = Objects.requireNonNull(originEntryType, "originEntryType is required"); if (!StringUtils.hasText(commandId)) { throw new IllegalArgumentException("commandId is required"); } @@ -175,6 +181,7 @@ public sealed interface GatewayCommand permits GatewayEnvelope envelope, ReplyMode replyMode, String commandId, + AgentRuntimeEntryType originEntryType, String sessionUid, int pageIndex, int pageSize) implements GatewayCommand { @@ -182,6 +189,8 @@ public sealed interface GatewayCommand permits public GetSessionTranscriptCommand { envelope = Objects.requireNonNull(envelope, "envelope is required"); replyMode = Objects.requireNonNull(replyMode, "replyMode is required"); + // The command must fully describe the session query scope. + originEntryType = Objects.requireNonNull(originEntryType, "originEntryType is required"); if (!StringUtils.hasText(commandId) || !StringUtils.hasText(sessionUid)) { throw new IllegalArgumentException("commandId and sessionUid are required"); } diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryService.java index 10fba771f3..b59e238cbb 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryService.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryService.java @@ -45,6 +45,7 @@ public class GatewayQueryService { public GatewaySingleResponse listSessions(ListSessionsCommand command) { Page sessions = sessionService.findSessions( command.envelope(), + command.originEntryType(), command.title(), PageRequest.of(command.pageIndex(), command.pageSize())); return GatewaySingleResponse.builder() @@ -59,7 +60,8 @@ public class GatewayQueryService { } public GatewaySingleResponse getSession(GetSessionCommand command) { - return sessionService.findOwnedSession(command.sessionUid(), command.envelope()) + return sessionService.findOwnedSession( + command.sessionUid(), command.envelope(), command.originEntryType()) .map(session -> GatewaySingleResponse.builder() .meta(Meta.builder() .commandId(command.commandId()) @@ -83,7 +85,8 @@ public class GatewayQueryService { public GatewaySingleResponse getSessionTranscript(GetSessionTranscriptCommand command) { PageRequest pageRequest = PageRequest.of(command.pageIndex(), command.pageSize()); - AgentSession session = sessionService.findOwnedSession(command.sessionUid(), command.envelope()).orElse(null); + AgentSession session = sessionService.findOwnedSession( + command.sessionUid(), command.envelope(), command.originEntryType()).orElse(null); Page transcript = session == null ? Page.empty(pageRequest) : sessionService.findTranscriptEntries(session.getId(), pageRequest); diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/core/ChannelId.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/core/ChannelId.java index 4f56aa546e..98e204d64c 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/core/ChannelId.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/core/ChannelId.java @@ -17,11 +17,11 @@ package org.apache.hertzbeat.ai.gateway.channel.core; -/** Built-in Agent Gateway channel identifiers. */ +/** Built-in Agent Gateway ingress channel identifiers. */ public enum ChannelId { WEB_UI("web-ui"), - ALERT("alert"); + SYSTEM("system"); private final String id; diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryController.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryController.java index 8820319b9c..a7dc95bd97 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryController.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryController.java @@ -32,6 +32,7 @@ import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.ReplyMode; import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.GatewaySingleResponse; import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; import org.apache.hertzbeat.ai.gateway.identity.AgentActor; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.common.entity.agent.AgentSession; import org.apache.hertzbeat.common.entity.agent.AgentTranscriptEntry; import org.apache.hertzbeat.common.entity.dto.Message; @@ -66,6 +67,7 @@ public class QueryController { .envelope(webUiEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("get-session:" + sessionId) + .originEntryType(AgentRuntimeEntryType.USER_INPUT) .sessionUid(sessionId) .build()))); } @@ -81,6 +83,7 @@ public class QueryController { .envelope(webUiEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("list-sessions:" + pageIndex) + .originEntryType(AgentRuntimeEntryType.USER_INPUT) .title(null) .pageIndex(pageIndex) .pageSize(pageSize) @@ -100,6 +103,7 @@ public class QueryController { .envelope(alertAnalysisEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("list-alert-analysis-sessions:" + pageIndex) + .originEntryType(AgentRuntimeEntryType.ALERT_TRIGGER) .title(search) .pageIndex(pageIndex) .pageSize(pageSize) @@ -116,6 +120,7 @@ public class QueryController { .envelope(alertAnalysisEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("get-alert-analysis-session:" + sessionId) + .originEntryType(AgentRuntimeEntryType.ALERT_TRIGGER) .sessionUid(sessionId) .build()))); } @@ -132,6 +137,7 @@ public class QueryController { .envelope(webUiEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("get-session-transcript:" + sessionUid) + .originEntryType(AgentRuntimeEntryType.USER_INPUT) .sessionUid(sessionUid) .pageIndex(pageIndex) .pageSize(pageSize) @@ -151,6 +157,7 @@ public class QueryController { .envelope(alertAnalysisEnvelope()) .replyMode(ReplyMode.FINAL_ONLY) .commandId("get-alert-analysis-session-transcript:" + sessionUid) + .originEntryType(AgentRuntimeEntryType.ALERT_TRIGGER) .sessionUid(sessionUid) .pageIndex(pageIndex) .pageSize(pageSize) @@ -169,7 +176,7 @@ public class QueryController { private GatewayEnvelope alertAnalysisEnvelope() { ActorSupport.requireCurrentAdminSurenessActor(); return GatewayEnvelope.builder() - .channelId(ChannelId.ALERT.id()) + .channelId(ChannelId.SYSTEM.id()) .receivedAt(System.currentTimeMillis()) .actor(AgentActor.alertAnalysisActor()) .build(); diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunService.java index da36f39167..0d492edee1 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunService.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunService.java @@ -19,10 +19,12 @@ package org.apache.hertzbeat.ai.gateway.conversation; import jakarta.persistence.EntityManager; import java.time.LocalDateTime; +import java.util.List; import java.util.Optional; import org.apache.hertzbeat.ai.gateway.conversation.persistence.AgentRunDao; import org.apache.hertzbeat.ai.gateway.contract.AgentTargetRef; import org.apache.hertzbeat.ai.gateway.contract.UserInput; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.ai.gateway.text.GatewayText; import org.apache.hertzbeat.common.entity.agent.AgentRun; import org.apache.hertzbeat.common.entity.agent.AgentSession; @@ -47,13 +49,14 @@ public class AgentRunService { this.entityManager = entityManager; } - public AgentRun createOrResumeRun(AgentSession session, UserInput userInput) { + public AgentRun createOrResumeRun(AgentSession session, UserInput userInput, + AgentRuntimeEntryType entryType) { String messageId = userInput.getMessageId(); Optional existingRun = runDao.findBySessionIdAndMessageId(session.getId(), messageId); if (existingRun.isPresent()) { return existingRun.get(); } - AgentRun run = buildRun(session, userInput, messageId); + AgentRun run = buildRun(session, userInput, messageId, entryType); try { return runDao.saveAndFlush(run); } catch (DataIntegrityViolationException e) { @@ -131,12 +134,29 @@ public class AgentRunService { .build(); } - private AgentRun buildRun(AgentSession session, UserInput userInput, String messageId) { + public Optional findCreatedRun(Long sessionId) { + return runDao.findFirstBySessionIdAndStatusOrderByGmtCreateAsc( + sessionId, AgentRunStatus.CREATED.name()); + } + + public Optional findRunningRun(Long sessionId) { + return runDao.findFirstBySessionIdAndStatusOrderByGmtCreateAsc( + sessionId, AgentRunStatus.RUNNING.name()); + } + + public boolean hasActiveRun(Long sessionId) { + return runDao.existsBySessionIdAndStatusIn(sessionId, + List.of(AgentRunStatus.CREATED.name(), AgentRunStatus.RUNNING.name())); + } + + private AgentRun buildRun(AgentSession session, UserInput userInput, String messageId, + AgentRuntimeEntryType entryType) { AgentTargetRef target = userInput.getTarget(); return AgentRun.builder() .runUid("run_" + SnowFlakeIdGenerator.generateId()) .sessionId(session.getId()) .messageId(messageId) + .entryType(entryType.name()) .targetMonitorId(target == null ? null : target.getMonitorId()) .targetAlertId(target == null ? null : target.getAlertId()) .targetCollector(target == null ? null : target.getCollector()) diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionService.java index 9b3e9abe1a..540573b454 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionService.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionService.java @@ -29,6 +29,7 @@ import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; import org.apache.hertzbeat.ai.gateway.contract.UserInput; import org.apache.hertzbeat.ai.gateway.contract.UserInput.Message; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeHistoryWindow; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.ai.gateway.runtime.TranscriptMessage; import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; import org.apache.hertzbeat.ai.gateway.text.GatewayText; @@ -67,7 +68,8 @@ public class AgentSessionService { this.entityManager = entityManager; } - public AgentSession findOrCreateSession(GatewayEnvelope envelope, UserInput userInput) { + public AgentSession findOrCreateSession(GatewayEnvelope envelope, UserInput userInput, + AgentRuntimeEntryType originEntryType) { AgentActor actor = envelope.getActor(); String sessionKey = sessionKeyBuilder.build(envelope, userInput.getConversationId()); Optional existed = sessionDao.findBySessionKey(sessionKey); @@ -79,6 +81,7 @@ public class AgentSessionService { .sessionUid("ags_" + SnowFlakeIdGenerator.generateId()) .sessionKey(sessionKey) .channel(envelope.getChannelId()) + .originEntryType(originEntryType.name()) .conversationId(userInput.getConversationId()) .actorType(actor.getType()) .actorId(actor.getId()) @@ -129,7 +132,8 @@ public class AgentSessionService { return sessionDao.findBySessionUid(normalized); } - public Optional findOwnedSession(String sessionId, GatewayEnvelope envelope) { + public Optional findOwnedSession( + String sessionId, GatewayEnvelope envelope, AgentRuntimeEntryType originEntryType) { AgentActor actor = envelope.getActor(); if (!ActorSupport.hasIdentity(actor)) { throw new IllegalArgumentException("Session query actor is required"); @@ -140,30 +144,41 @@ public class AgentSessionService { return Optional.empty(); } if (normalized.chars().allMatch(Character::isDigit)) { - return sessionDao.findByIdAndChannelAndActorTypeAndActorId( - Long.parseLong(normalized), envelope.getChannelId(), actor.getType(), actor.getId()); + return sessionDao.findByIdAndChannelAndActorTypeAndActorIdAndOriginEntryType( + Long.parseLong(normalized), envelope.getChannelId(), actor.getType(), actor.getId(), + originEntryType.name()); } - return sessionDao.findBySessionUidAndChannelAndActorTypeAndActorId( - normalized, envelope.getChannelId(), actor.getType(), actor.getId()); + return sessionDao.findBySessionUidAndChannelAndActorTypeAndActorIdAndOriginEntryType( + normalized, envelope.getChannelId(), actor.getType(), actor.getId(), originEntryType.name()); } - public Page findSessions(GatewayEnvelope envelope, String title, Pageable pageable) { + public Page findSessions( + GatewayEnvelope envelope, AgentRuntimeEntryType originEntryType, String title, Pageable pageable) { AgentActor actor = envelope.getActor(); if (!ActorSupport.hasIdentity(actor)) { throw new IllegalArgumentException("Session query actor is required"); } if (!StringUtils.hasText(title)) { - return sessionDao.findByChannelAndActorTypeAndActorIdOrderByGmtUpdateDesc( - envelope.getChannelId(), actor.getType(), actor.getId(), pageable); + return sessionDao.findByChannelAndActorTypeAndActorIdAndOriginEntryTypeOrderByGmtUpdateDesc( + envelope.getChannelId(), actor.getType(), actor.getId(), originEntryType.name(), pageable); } - return sessionDao.findByChannelAndActorTypeAndActorIdAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( - envelope.getChannelId(), actor.getType(), actor.getId(), title, pageable); + return sessionDao + .findByChannelAndActorTypeAndActorIdAndOriginEntryTypeAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( + envelope.getChannelId(), actor.getType(), actor.getId(), originEntryType.name(), title, pageable); } public Page findTranscriptEntries(Long sessionId, Pageable pageable) { return transcriptEntryDao.findBySessionIdOrderBySessionSequenceAsc(sessionId, pageable); } + public Page findConversationTranscriptEntries(Long sessionId, Pageable pageable) { + return transcriptEntryDao.findBySessionIdAndMessageRoleInOrderBySessionSequenceDesc( + sessionId, + List.of(TranscriptMessage.TranscriptRole.USER.wireValue(), + TranscriptMessage.TranscriptRole.ASSISTANT.wireValue()), + pageable); + } + @Transactional public List findRecentTranscriptMessages(Long sessionId) { if (sessionId == null) { diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentRunDao.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentRunDao.java index 3071f6249b..e733863f4e 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentRunDao.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentRunDao.java @@ -17,6 +17,7 @@ package org.apache.hertzbeat.ai.gateway.conversation.persistence; +import java.util.List; import java.util.Optional; import org.apache.hertzbeat.common.entity.agent.AgentRun; import org.springframework.data.jpa.repository.JpaRepository; @@ -37,4 +38,9 @@ public interface AgentRunDao extends JpaRepository { * Find the existing run for a session event. */ Optional findBySessionIdAndMessageId(Long sessionId, String messageId); + + Optional findFirstBySessionIdAndStatusOrderByGmtCreateAsc(Long sessionId, String status); + + boolean existsBySessionIdAndStatusIn(Long sessionId, List statuses); + } diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentSessionDao.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentSessionDao.java index 1aa394bbe9..6e93e556cb 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentSessionDao.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentSessionDao.java @@ -42,17 +42,19 @@ public interface AgentSessionDao extends JpaRepository { */ Optional findBySessionUid(String sessionUid); - Optional findByIdAndChannelAndActorTypeAndActorId( - Long id, String channel, String actorType, String actorId); + Optional findByIdAndChannelAndActorTypeAndActorIdAndOriginEntryType( + Long id, String channel, String actorType, String actorId, String originEntryType); - Optional findBySessionUidAndChannelAndActorTypeAndActorId( - String sessionUid, String channel, String actorType, String actorId); + Optional findBySessionUidAndChannelAndActorTypeAndActorIdAndOriginEntryType( + String sessionUid, String channel, String actorType, String actorId, String originEntryType); - Page findByChannelAndActorTypeAndActorIdOrderByGmtUpdateDesc( - String channel, String actorType, String actorId, Pageable pageable); + Page findByChannelAndActorTypeAndActorIdAndOriginEntryTypeOrderByGmtUpdateDesc( + String channel, String actorType, String actorId, String originEntryType, Pageable pageable); - Page findByChannelAndActorTypeAndActorIdAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( - String channel, String actorType, String actorId, String title, Pageable pageable); + Page + findByChannelAndActorTypeAndActorIdAndOriginEntryTypeAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( + String channel, String actorType, String actorId, String originEntryType, + String title, Pageable pageable); /** * Lock a session row while assigning transcript append sequence. diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentTranscriptEntryDao.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentTranscriptEntryDao.java index 360dece09a..8270f90342 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentTranscriptEntryDao.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/conversation/persistence/AgentTranscriptEntryDao.java @@ -37,6 +37,12 @@ public interface AgentTranscriptEntryDao extends JpaRepository findBySessionIdOrderBySessionSequenceAsc(Long sessionId, Pageable pageable); + /** + * Find conversation messages for a session from newest to oldest. + */ + Page findBySessionIdAndMessageRoleInOrderBySessionSequenceDesc( + Long sessionId, List messageRoles, Pageable pageable); + /** * Find recent transcript entries for a session in reverse append order. */ diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/ActorSupport.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/ActorSupport.java index c0bf0b8556..b563a7fc47 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/ActorSupport.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/ActorSupport.java @@ -38,6 +38,8 @@ public final class ActorSupport { public static final String ID_ALERT_ANALYSIS = "alert-analysis"; + public static final String ID_SCHEDULE = "schedule"; + public static final String ROLE_ADMIN = "admin"; public static final String ROLE_USER = "user"; @@ -46,6 +48,8 @@ public final class ActorSupport { public static final String ROLE_ALERT_ANALYSIS = "alert-analysis"; + public static final String ROLE_SCHEDULE = "schedule"; + private static final String ROLES_CLAIM = "roles"; private static final List WELL_KNOWN_ROLES = List.of(ROLE_ADMIN, ROLE_USER, ROLE_GUEST); diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/AgentActor.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/AgentActor.java index 880dc98afb..9393ff316f 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/AgentActor.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/identity/AgentActor.java @@ -48,4 +48,12 @@ public class AgentActor { .roles(List.of(ActorSupport.ROLE_ALERT_ANALYSIS)) .build(); } + + public static AgentActor scheduleActor() { + return AgentActor.builder() + .type(ActorSupport.TYPE_SYSTEM) + .id(ActorSupport.ID_SCHEDULE) + .roles(List.of(ActorSupport.ROLE_SCHEDULE)) + .build(); + } } diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilder.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilder.java index ea4c11afa6..e7ac06707c 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilder.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilder.java @@ -19,6 +19,7 @@ package org.apache.hertzbeat.ai.gateway.runtime; import java.time.Clock; import java.time.Instant; +import java.time.ZoneId; import java.time.format.DateTimeFormatter; import java.util.List; import java.util.Objects; @@ -60,8 +61,9 @@ public class AgentRuntimeContextBuilder { AgentTargetRef effectiveTarget = effectiveTarget(entryType, userInput, run); List chatHistory = List.copyOf(request.getChatHistory()); Instant now = Instant.now(clock); - String currentTimeIso = DateTimeFormatter.ISO_OFFSET_DATE_TIME.format(now.atZone(clock.getZone())); - String timezone = clock.getZone().getId(); + ZoneId systemZone = ZoneId.systemDefault(); + String currentTimeIso = DateTimeFormatter.ISO_OFFSET_DATE_TIME.format(now.atZone(systemZone)); + String timezone = systemZone.getId(); String traceId = resolveTraceId(); return AgentRuntimeContext.builder() .entryType(entryType) diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommand.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentSchedule.java similarity index 61% rename from hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommand.java rename to hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentSchedule.java index 13f80995a6..e31a60a176 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommand.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentSchedule.java @@ -18,6 +18,7 @@ package org.apache.hertzbeat.ai.gateway.schedule; import jakarta.persistence.Column; +import jakarta.persistence.Convert; import jakarta.persistence.Entity; import jakarta.persistence.EntityListeners; import jakarta.persistence.GeneratedValue; @@ -26,53 +27,42 @@ import jakarta.persistence.Id; import jakarta.persistence.Index; import jakarta.persistence.Table; import java.time.LocalDateTime; +import java.util.List; import lombok.AllArgsConstructor; import lombok.Builder; import lombok.Data; import lombok.NoArgsConstructor; +import org.apache.hertzbeat.common.entity.manager.JsonLongListAttributeConverter; +import org.springframework.data.annotation.CreatedBy; import org.springframework.data.annotation.CreatedDate; +import org.springframework.data.annotation.LastModifiedBy; import org.springframework.data.annotation.LastModifiedDate; import org.springframework.data.jpa.domain.support.AuditingEntityListener; /** - * Recurring Gateway command that replays a user input through the standard runtime entry. + * System-level recurring Agent inspection rule. */ @Data @Builder @Entity @EntityListeners(AuditingEntityListener.class) -@Table(name = "hzb_agent_scheduled_command", indexes = { - @Index(name = "idx_agent_scheduled_command_session", columnList = "session_id"), - @Index(name = "idx_agent_scheduled_command_due", columnList = "enabled, next_run_time") +@Table(name = "hzb_agent_schedule", indexes = { + @Index(name = "idx_agent_schedule_due", columnList = "enabled, next_trigger_at"), + @Index(name = "idx_agent_schedule_session", columnList = "session_id") }) @AllArgsConstructor @NoArgsConstructor -public class AgentScheduledCommand { +public class AgentSchedule { @Id @GeneratedValue(strategy = GenerationType.IDENTITY) private Long id; - @Column(name = "session_id", nullable = false) - private Long sessionId; + @Column(nullable = false, length = 128) + private String name; - @Column(name = "channel", nullable = false, length = 64) - private String channel; - - @Column(name = "conversation_id", nullable = false, length = 256) - private String conversationId; - - @Column(name = "actor_type", nullable = false, length = 64) - private String actorType; - - @Column(name = "actor_id", nullable = false, length = 128) - private String actorId; - - @Column(name = "actor_roles", nullable = false, length = 1024) - private String actorRoles; - - @Column(name = "message", nullable = false, length = 4096) - private String message; + @Column(nullable = false, length = 4096) + private String instruction; @Column(name = "cron_expression", nullable = false, length = 64) private String cronExpression; @@ -81,11 +71,32 @@ public class AgentScheduledCommand { @Column(nullable = false) private boolean enabled = true; - @Column(name = "last_run_time") - private LocalDateTime lastRunTime; + @Column(name = "session_id") + private Long sessionId; - @Column(name = "next_run_time") - private LocalDateTime nextRunTime; + @Convert(converter = JsonLongListAttributeConverter.class) + @Column(name = "receiver_ids", nullable = false, length = 2048) + private List receiverIds; + + @Column(name = "template_id") + private Long templateId; + + @Column(name = "created_from_session_uid", length = 64) + private String createdFromSessionUid; + + @Column(name = "last_trigger_at") + private Long lastTriggerAt; + + @Column(name = "next_trigger_at") + private Long nextTriggerAt; + + @CreatedBy + @Column(length = 64) + private String creator; + + @LastModifiedBy + @Column(length = 64) + private String modifier; @CreatedDate @Column(name = "gmt_create") diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleController.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleController.java new file mode 100644 index 0000000000..26160ab822 --- /dev/null +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleController.java @@ -0,0 +1,144 @@ +/* + * 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.gateway.schedule; + +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.tags.Tag; +import java.util.List; +import org.apache.hertzbeat.ai.gateway.conversation.AgentSessionService; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.apache.hertzbeat.common.entity.agent.AgentTranscriptEntry; +import org.apache.hertzbeat.common.entity.dto.Message; +import org.springframework.data.domain.Page; +import org.springframework.data.domain.PageRequest; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.DeleteMapping; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.PatchMapping; +import org.springframework.web.bind.annotation.PathVariable; +import org.springframework.web.bind.annotation.PostMapping; +import org.springframework.web.bind.annotation.PutMapping; +import org.springframework.web.bind.annotation.RequestBody; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RequestParam; +import org.springframework.web.bind.annotation.RestController; + +/** + * System-level Agent schedule API. + */ +@Tag(name = "Agent Schedule API") +@RestController +@RequestMapping("/api/agent/schedules") +public class AgentScheduleController { + + private final AgentScheduleService scheduleService; + private final AgentScheduleExecutor scheduleExecutor; + private final AgentSessionService sessionService; + + public AgentScheduleController(AgentScheduleService scheduleService, + AgentScheduleExecutor scheduleExecutor, + AgentSessionService sessionService) { + this.scheduleService = scheduleService; + this.scheduleExecutor = scheduleExecutor; + this.sessionService = sessionService; + } + + @PostMapping + @Operation(summary = "Create an Agent schedule") + public ResponseEntity> create(@RequestBody ScheduleRequest request) { + return ResponseEntity.ok(Message.success(scheduleService.create(request.toEntity()))); + } + + @PutMapping("/{scheduleId}") + @Operation(summary = "Update an Agent schedule") + public ResponseEntity> update(@PathVariable Long scheduleId, + @RequestBody ScheduleRequest request) { + return ResponseEntity.ok(Message.success(scheduleService.update(scheduleId, request.toEntity()))); + } + + @GetMapping + @Operation(summary = "List Agent schedules") + public ResponseEntity>> list( + @RequestParam(defaultValue = "0") int pageIndex, + @RequestParam(defaultValue = "20") int pageSize) { + return ResponseEntity.ok(Message.success( + scheduleService.list(PageRequest.of(pageIndex, pageSize)))); + } + + @GetMapping("/{scheduleId}") + @Operation(summary = "Get an Agent schedule") + public ResponseEntity> get(@PathVariable Long scheduleId) { + return ResponseEntity.ok(Message.success(scheduleService.get(scheduleId))); + } + + @DeleteMapping("/{scheduleId}") + @Operation(summary = "Delete an Agent schedule") + public ResponseEntity> delete(@PathVariable Long scheduleId) { + scheduleService.delete(scheduleId); + return ResponseEntity.ok(Message.success("Agent schedule deleted")); + } + + @PatchMapping("/{scheduleId}/enabled") + @Operation(summary = "Enable or disable an Agent schedule") + public ResponseEntity> toggle(@PathVariable Long scheduleId, + @RequestParam boolean enabled) { + return ResponseEntity.ok(Message.success(scheduleService.toggle(scheduleId, enabled))); + } + + @PostMapping("/{scheduleId}/run") + @Operation(summary = "Run an Agent schedule immediately") + public ResponseEntity> runNow(@PathVariable Long scheduleId) { + return ResponseEntity.ok(Message.success(scheduleExecutor.executeNow(scheduleId))); + } + + @GetMapping("/{scheduleId}/transcript") + @Operation(summary = "List the fixed Agent session transcript for a schedule") + public ResponseEntity>> transcript( + @PathVariable Long scheduleId, + @RequestParam(defaultValue = "0") int pageIndex, + @RequestParam(defaultValue = "20") int pageSize) { + AgentSchedule schedule = scheduleService.get(scheduleId); + PageRequest pageable = PageRequest.of(pageIndex, pageSize); + Page transcript = schedule.getSessionId() == null + ? Page.empty(pageable) + : sessionService.findConversationTranscriptEntries(schedule.getSessionId(), pageable); + return ResponseEntity.ok(Message.success(transcript)); + } + + /** + * Create and update boundary for schedule-owned fields. + */ + public record ScheduleRequest( + String name, + String instruction, + String cronExpression, + Boolean enabled, + List receiverIds, + Long templateId) { + + AgentSchedule toEntity() { + return AgentSchedule.builder() + .name(name) + .instruction(instruction) + .cronExpression(cronExpression) + .enabled(enabled == null || enabled) + .receiverIds(receiverIds) + .templateId(templateId) + .build(); + } + } +} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandDao.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleDao.java similarity index 52% rename from hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandDao.java rename to hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleDao.java index 4d4e3c4302..4802013622 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandDao.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleDao.java @@ -7,7 +7,6 @@ * 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. @@ -17,18 +16,33 @@ package org.apache.hertzbeat.ai.gateway.schedule; -import java.time.LocalDateTime; import java.util.List; +import java.util.Optional; +import org.springframework.data.domain.Pageable; import org.springframework.data.jpa.repository.JpaRepository; +import org.springframework.data.jpa.repository.Query; +import org.springframework.data.repository.query.Param; import org.springframework.stereotype.Repository; /** - * Persistence for scheduled Gateway commands. + * Agent schedule persistence. */ @Repository -public interface AgentScheduledCommandDao extends JpaRepository { +public interface AgentScheduleDao extends JpaRepository { - List findBySessionIdOrderByIdAsc(Long sessionId); + List findByEnabledTrueAndNextTriggerAtLessThanEqualOrderByNextTriggerAtAsc( + Long now, Pageable pageable); - List findByEnabledTrueAndNextRunTimeLessThanEqual(LocalDateTime now); + List findByEnabledTrue(); + + Optional findBySessionId(Long sessionId); + + @Query(""" + select schedule from AgentSchedule schedule + where schedule.sessionId in ( + select run.sessionId from AgentRun run where run.status = :status + ) + order by schedule.gmtUpdate asc + """) + List findWithRuns(@Param("status") String status, Pageable pageable); } diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutor.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutor.java new file mode 100644 index 0000000000..dd627c42f5 --- /dev/null +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutor.java @@ -0,0 +1,190 @@ +/* + * 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.gateway.schedule; + +import jakarta.annotation.PreDestroy; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ArrayBlockingQueue; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import lombok.extern.slf4j.Slf4j; +import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.InvokeCommand; +import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.ReplyMode; +import org.apache.hertzbeat.ai.gateway.application.GatewayCommandRouter; +import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.GatewaySingleResponse; +import org.apache.hertzbeat.ai.gateway.channel.core.ChannelId; +import org.apache.hertzbeat.ai.gateway.contract.AgentResponseLanguage; +import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; +import org.apache.hertzbeat.ai.gateway.contract.UserInput; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunService; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunStatus; +import org.apache.hertzbeat.ai.gateway.identity.AgentActor; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.springframework.boot.context.event.ApplicationReadyEvent; +import org.springframework.context.event.EventListener; +import org.springframework.scheduling.annotation.Scheduled; +import org.springframework.stereotype.Component; +import org.springframework.util.StringUtils; + +/** + * Claims due schedules and executes their standard Gateway commands asynchronously. + */ +@Slf4j +@Component +public class AgentScheduleExecutor { + + private final AgentScheduleService scheduleService; + private final AgentRunService runService; + private final GatewayCommandRouter commandRouter; + private final AgentScheduleNoticeService noticeService; + private final Set submittedSchedules = ConcurrentHashMap.newKeySet(); + private final ExecutorService executor = new ThreadPoolExecutor(4, 4, 0, TimeUnit.MILLISECONDS, + new ArrayBlockingQueue<>(256), Thread.ofPlatform().name("agent-schedule-", 0).factory(), + new ThreadPoolExecutor.AbortPolicy()); + + public AgentScheduleExecutor(AgentScheduleService scheduleService, + AgentRunService runService, + GatewayCommandRouter commandRouter, + AgentScheduleNoticeService noticeService) { + this.scheduleService = scheduleService; + this.runService = runService; + this.commandRouter = commandRouter; + this.noticeService = noticeService; + } + + @EventListener(ApplicationReadyEvent.class) + public void failInterruptedRuns() { + for (AgentSchedule schedule : scheduleService.findInterrupted()) { + runService.findRunningRun(schedule.getSessionId()).ifPresent(run -> { + String message = "Agent schedule execution was interrupted by process restart"; + AgentRun failed = runService.markFailed(run, message); + noticeService.send(schedule, failed, false, message); + }); + } + } + + @Scheduled(fixedDelay = 10_000L) + public void executeDueSchedules() { + for (AgentSchedule schedule : scheduleService.findPending()) { + if (!submittedSchedules.add(schedule.getId())) { + continue; + } + try { + scheduleService.findCreatedRun(schedule) + .ifPresentOrElse(run -> submit(schedule, run), + () -> submittedSchedules.remove(schedule.getId())); + } catch (RuntimeException exception) { + submittedSchedules.remove(schedule.getId()); + log.error("Failed to recover pending Agent schedule {}", schedule.getId(), exception); + } + } + long now = System.currentTimeMillis(); + for (AgentSchedule schedule : scheduleService.findDue(now)) { + if (!submittedSchedules.add(schedule.getId())) { + continue; + } + try { + scheduleService.claimCronRun(schedule.getId(), now) + .ifPresentOrElse(run -> submit(scheduleService.get(schedule.getId()), run), + () -> submittedSchedules.remove(schedule.getId())); + } catch (RuntimeException exception) { + submittedSchedules.remove(schedule.getId()); + log.error("Failed to claim due Agent schedule {}", schedule.getId(), exception); + } + } + } + + public AgentRun executeNow(Long scheduleId) { + if (!submittedSchedules.add(scheduleId)) { + throw new IllegalStateException("Agent schedule already has an active run"); + } + try { + AgentRun run = scheduleService.claimManualRun(scheduleId); + submit(scheduleService.get(scheduleId), run); + return run; + } catch (RuntimeException exception) { + submittedSchedules.remove(scheduleId); + throw exception; + } + } + + private void submit(AgentSchedule schedule, AgentRun run) { + try { + executor.execute(() -> execute(schedule, run)); + } catch (RejectedExecutionException exception) { + submittedSchedules.remove(schedule.getId()); + log.warn("Agent schedule {} execution was rejected: {}", + schedule.getId(), exception.getMessage()); + } + } + + private void execute(AgentSchedule schedule, AgentRun run) { + try { + GatewaySingleResponse response = (GatewaySingleResponse) commandRouter.handle(command(schedule, run)); + Map body = response.body() instanceof Map map ? map : Map.of(); + boolean succeeded = AgentRunStatus.SUCCEEDED.name().equals(body.get("status")); + Object responseMessage = body.get("message"); + String message = responseMessage == null + ? (succeeded ? "Agent schedule completed" : "Agent schedule failed") + : String.valueOf(responseMessage); + noticeService.send(schedule, runService.findRun(run.getRunUid()).orElse(run), succeeded, message); + } catch (RuntimeException exception) { + AgentRun current = runService.findRun(run.getRunUid()).orElse(run); + // Runtime failures such as interrupted providers may not carry a message; persist a useful terminal reason. + String failureMessage = StringUtils.hasText(exception.getMessage()) + ? exception.getMessage() + : "Agent schedule execution failed"; + if (!AgentRunStatus.FAILED.name().equals(current.getStatus())) { + current = runService.markFailed(current, failureMessage); + } + noticeService.send(schedule, current, false, failureMessage); + log.error("Agent schedule {} run {} failed", schedule.getId(), run.getRunUid(), exception); + } finally { + submittedSchedules.remove(schedule.getId()); + } + } + + private InvokeCommand command(AgentSchedule schedule, AgentRun run) { + long now = System.currentTimeMillis(); + return InvokeCommand.builder() + .envelope(GatewayEnvelope.builder() + .channelId(ChannelId.SYSTEM.id()) + .receivedAt(now) + .preferredLanguage(AgentResponseLanguage.systemDefault()) + .actor(AgentActor.scheduleActor()) + .build()) + .replyMode(ReplyMode.FINAL_ONLY) + .commandId(run.getMessageId()) + .entryType(AgentRuntimeEntryType.SCHEDULE_TRIGGER) + .userInput(UserInput.builder() + .messageId(run.getMessageId()) + .conversationId("schedule:" + schedule.getId()) + .message(UserInput.Message.builder().text(schedule.getInstruction()).build()) + .build()) + .build(); + } + + @PreDestroy + public void close() { + executor.shutdownNow(); + } +} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeService.java new file mode 100644 index 0000000000..415a7cce69 --- /dev/null +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeService.java @@ -0,0 +1,122 @@ +/* + * 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.gateway.schedule; + +import java.time.ZoneId; +import java.util.List; +import java.util.Map; +import java.util.concurrent.RejectedExecutionException; +import lombok.extern.slf4j.Slf4j; +import org.apache.hertzbeat.alert.AlerterWorkerPool; +import org.apache.hertzbeat.alert.notice.AlertNoticeDispatch; +import org.apache.hertzbeat.alert.service.NoticeConfigService; +import org.apache.hertzbeat.ai.gateway.text.GatewayText; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.apache.hertzbeat.common.entity.alerter.GroupAlert; +import org.apache.hertzbeat.common.entity.alerter.NoticeReceiver; +import org.apache.hertzbeat.common.entity.alerter.NoticeTemplate; +import org.apache.hertzbeat.common.entity.alerter.SingleAlert; +import org.springframework.stereotype.Service; +import org.springframework.util.StringUtils; + +/** + * Sends schedule results through the existing alert notification transports. + */ +@Slf4j +@Service +public class AgentScheduleNoticeService { + + private final NoticeConfigService noticeConfigService; + private final AlertNoticeDispatch alertNoticeDispatch; + private final AlerterWorkerPool workerPool; + + public AgentScheduleNoticeService(NoticeConfigService noticeConfigService, + AlertNoticeDispatch alertNoticeDispatch, + AlerterWorkerPool workerPool) { + this.noticeConfigService = noticeConfigService; + this.alertNoticeDispatch = alertNoticeDispatch; + this.workerPool = workerPool; + } + + public void send(AgentSchedule schedule, AgentRun run, boolean succeeded, String result) { + NoticeTemplate template = schedule.getTemplateId() == null + ? null + : noticeConfigService.getOneTemplateById(schedule.getTemplateId()); + GroupAlert alert = scheduleAlert(schedule, run, succeeded, result); + for (Long receiverId : schedule.getReceiverIds()) { + NoticeReceiver receiver = noticeConfigService.getReceiverById(receiverId); + if (receiver == null || receiver.getType() == null) { + log.warn("Agent schedule {} skipped missing notice receiver {}", schedule.getId(), receiverId); + continue; + } + try { + workerPool.executeNotify(receiver.getType(), () -> { + try { + alertNoticeDispatch.sendNoticeMsg(receiver, template, alert); + } catch (RuntimeException exception) { + log.warn("Agent schedule {} failed to notify receiver {}: {}", + schedule.getId(), receiverId, exception.getMessage()); + } + }); + } catch (RejectedExecutionException exception) { + log.warn("Agent schedule {} notification was rejected for receiver {}: {}", + schedule.getId(), receiverId, exception.getMessage()); + } + } + } + + private GroupAlert scheduleAlert(AgentSchedule schedule, AgentRun run, boolean succeeded, String result) { + String status = succeeded ? "resolved" : "firing"; + String severity = succeeded ? "info" : "critical"; + // Model output crosses into external notification channels, so bound it to the existing alert content limit. + String content = GatewayText.safeSummary( + StringUtils.hasText(result) ? result : "Agent schedule execution failed", 4096); + long startAt = run.getStartedAt() == null + ? System.currentTimeMillis() + : run.getStartedAt().atZone(ZoneId.systemDefault()).toInstant().toEpochMilli(); + long endAt = run.getCompletedAt() == null + ? System.currentTimeMillis() + : run.getCompletedAt().atZone(ZoneId.systemDefault()).toInstant().toEpochMilli(); + Map labels = Map.of( + "alertname", "Agent Schedule: " + schedule.getName(), + "severity", severity, + "source", "agent_schedule", + "scheduleId", String.valueOf(schedule.getId())); + SingleAlert singleAlert = SingleAlert.builder() + .status(status) + .labels(labels) + .annotations(succeeded ? Map.of() : Map.of("error", content)) + .content(content) + .triggerTimes(1) + .startAt(startAt) + .activeAt(startAt) + .endAt(endAt) + .build(); + return GroupAlert.builder() + .status(status) + .groupLabels(Map.of("source", "agent_schedule", + "scheduleId", String.valueOf(schedule.getId()))) + .commonLabels(labels) + .commonAnnotations(Map.of( + "scheduleName", schedule.getName(), + "runUid", run.getRunUid(), + "resultStatus", succeeded ? "SUCCEEDED" : "FAILED", + "triggeredAt", String.valueOf(schedule.getLastTriggerAt()))) + .alerts(List.of(singleAlert)) + .build(); + } +} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleService.java new file mode 100644 index 0000000000..35133745d9 --- /dev/null +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleService.java @@ -0,0 +1,288 @@ +/* + * 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.gateway.schedule; + +import java.time.ZoneId; +import java.time.ZonedDateTime; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Optional; +import org.apache.hertzbeat.alert.service.NoticeConfigService; +import org.apache.hertzbeat.ai.gateway.channel.core.ChannelId; +import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; +import org.apache.hertzbeat.ai.gateway.contract.UserInput; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunService; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunStatus; +import org.apache.hertzbeat.ai.gateway.conversation.AgentSessionService; +import org.apache.hertzbeat.ai.gateway.identity.AgentActor; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.apache.hertzbeat.common.entity.agent.AgentSession; +import org.apache.hertzbeat.common.entity.alerter.NoticeReceiver; +import org.apache.hertzbeat.common.entity.alerter.NoticeTemplate; +import org.apache.hertzbeat.common.util.SnowFlakeIdGenerator; +import org.springframework.context.event.EventListener; +import org.springframework.data.domain.Page; +import org.springframework.data.domain.PageRequest; +import org.springframework.data.domain.Pageable; +import org.springframework.scheduling.support.CronExpression; +import org.springframework.stereotype.Service; +import org.springframework.transaction.annotation.Transactional; +import org.springframework.util.StringUtils; +import org.apache.hertzbeat.common.support.event.SystemConfigChangeEvent; + +/** + * Agent schedule lifecycle and trigger claiming. + */ +@Service +public class AgentScheduleService { + + private static final int DISPATCH_BATCH_SIZE = 100; + + private final AgentScheduleDao scheduleDao; + private final AgentSessionService sessionService; + private final AgentRunService runService; + private final NoticeConfigService noticeConfigService; + + public AgentScheduleService(AgentScheduleDao scheduleDao, + AgentSessionService sessionService, + AgentRunService runService, + NoticeConfigService noticeConfigService) { + this.scheduleDao = scheduleDao; + this.sessionService = sessionService; + this.runService = runService; + this.noticeConfigService = noticeConfigService; + } + + @Transactional + public AgentSchedule create(AgentSchedule schedule) { + validate(schedule); + schedule.setId(null); + schedule.setSessionId(null); + schedule.setLastTriggerAt(null); + schedule.setNextTriggerAt(schedule.isEnabled() ? nextTriggerAt(schedule.getCronExpression()) : null); + AgentSchedule saved = scheduleDao.saveAndFlush(schedule); + saved.setSessionId(scheduleSession(saved).getId()); + return scheduleDao.save(saved); + } + + @Transactional + public AgentSchedule update(Long scheduleId, AgentSchedule input) { + AgentSchedule schedule = get(scheduleId); + rejectWhileRunning(schedule); + schedule.setName(input.getName()); + schedule.setInstruction(input.getInstruction()); + schedule.setCronExpression(input.getCronExpression()); + schedule.setReceiverIds(input.getReceiverIds()); + schedule.setTemplateId(input.getTemplateId()); + validate(schedule); + schedule.setNextTriggerAt(schedule.isEnabled() ? nextTriggerAt(schedule.getCronExpression()) : null); + return scheduleDao.save(schedule); + } + + public AgentSchedule get(Long scheduleId) { + return scheduleDao.findById(scheduleId) + .orElseThrow(() -> new IllegalArgumentException("Agent schedule not found: " + scheduleId)); + } + + public Page list(Pageable pageable) { + return scheduleDao.findAll(pageable); + } + + @Transactional + public AgentSchedule toggle(Long scheduleId, boolean enabled) { + AgentSchedule schedule = get(scheduleId); + if (!enabled) { + rejectWhileRunning(schedule); + } else { + validate(schedule); + } + schedule.setEnabled(enabled); + schedule.setNextTriggerAt(enabled ? nextTriggerAt(schedule.getCronExpression()) : null); + return scheduleDao.save(schedule); + } + + @Transactional + public void delete(Long scheduleId) { + AgentSchedule schedule = get(scheduleId); + rejectWhileRunning(schedule); + scheduleDao.delete(schedule); + } + + public List findDue(long now) { + return scheduleDao.findByEnabledTrueAndNextTriggerAtLessThanEqualOrderByNextTriggerAtAsc( + now, PageRequest.of(0, DISPATCH_BATCH_SIZE)); + } + + public List findPending() { + return scheduleDao.findWithRuns(AgentRunStatus.CREATED.name(), + PageRequest.of(0, DISPATCH_BATCH_SIZE)); + } + + public List findInterrupted() { + return scheduleDao.findWithRuns(AgentRunStatus.RUNNING.name(), + Pageable.unpaged()); + } + + @Transactional + public Optional claimCronRun(Long scheduleId, long now) { + AgentSchedule schedule = get(scheduleId); + if (!schedule.isEnabled() || schedule.getNextTriggerAt() == null || schedule.getNextTriggerAt() > now) { + return Optional.empty(); + } + long plannedAt = schedule.getNextTriggerAt(); + schedule.setLastTriggerAt(plannedAt); + schedule.setNextTriggerAt(nextTriggerAt(schedule.getCronExpression())); + if (runService.hasActiveRun(schedule.getSessionId())) { + scheduleDao.save(schedule); + return Optional.empty(); + } + AgentRun run = createRun(schedule, "schedule:" + schedule.getId() + ":cron:" + plannedAt); + scheduleDao.save(schedule); + return Optional.of(run); + } + + @Transactional + public AgentRun claimManualRun(Long scheduleId) { + AgentSchedule schedule = get(scheduleId); + if (runService.hasActiveRun(schedule.getSessionId())) { + throw new IllegalStateException("Agent schedule already has an active run"); + } + schedule.setLastTriggerAt(System.currentTimeMillis()); + AgentRun run = createRun(schedule, + "schedule:" + schedule.getId() + ":manual:" + SnowFlakeIdGenerator.generateId()); + scheduleDao.save(schedule); + return run; + } + + public Optional findCreatedRun(AgentSchedule schedule) { + return runService.findCreatedRun(schedule.getSessionId()); + } + + public Optional findBySessionId(Long sessionId) { + return scheduleDao.findBySessionId(sessionId); + } + + @EventListener(SystemConfigChangeEvent.class) + @Transactional + public void onSystemConfigChanged(SystemConfigChangeEvent event) { + for (AgentSchedule schedule : scheduleDao.findByEnabledTrue()) { + schedule.setNextTriggerAt(nextTriggerAt(schedule.getCronExpression())); + } + } + + private AgentRun createRun(AgentSchedule schedule, String messageId) { + AgentSession session = scheduleSession(schedule); + UserInput input = scheduleInput(schedule, messageId); + return runService.createOrResumeRun(session, input, AgentRuntimeEntryType.SCHEDULE_TRIGGER); + } + + private AgentSession scheduleSession(AgentSchedule schedule) { + if (schedule.getSessionId() != null) { + return sessionService.findSession(String.valueOf(schedule.getSessionId())) + .orElseThrow(() -> new IllegalStateException( + "Agent schedule session not found: " + schedule.getSessionId())); + } + long now = System.currentTimeMillis(); + return sessionService.findOrCreateSession( + GatewayEnvelope.builder() + .channelId(ChannelId.SYSTEM.id()) + .receivedAt(now) + .actor(AgentActor.scheduleActor()) + .build(), + scheduleInput(schedule, "schedule-session:" + schedule.getId()), + AgentRuntimeEntryType.SCHEDULE_TRIGGER); + } + + private UserInput scheduleInput(AgentSchedule schedule, String messageId) { + return UserInput.builder() + .messageId(messageId) + .conversationId("schedule:" + schedule.getId()) + .message(UserInput.Message.builder().text(schedule.getInstruction()).build()) + .build(); + } + + private void rejectWhileRunning(AgentSchedule schedule) { + if (schedule.getSessionId() != null && runService.hasActiveRun(schedule.getSessionId())) { + throw new IllegalStateException("Agent schedule cannot be changed while a run is active"); + } + } + + private void validate(AgentSchedule schedule) { + if (!StringUtils.hasText(schedule.getName()) || schedule.getName().length() > 128) { + throw new IllegalArgumentException("Agent schedule name is required and must not exceed 128 characters"); + } + if (!StringUtils.hasText(schedule.getInstruction()) || schedule.getInstruction().length() > 4096) { + throw new IllegalArgumentException( + "Agent schedule instruction is required and must not exceed 4096 characters"); + } + if (!StringUtils.hasText(schedule.getCronExpression()) || schedule.getCronExpression().length() > 64) { + throw new IllegalArgumentException( + "Agent schedule cron expression is required and must not exceed 64 characters"); + } + // API and tool callers may repeat receiver IDs; this boundary prevents duplicate notifications. + List receiverIds = schedule.getReceiverIds() == null + ? List.of() + : List.copyOf(new LinkedHashSet<>(schedule.getReceiverIds())); + if (receiverIds.isEmpty() || receiverIds.stream().anyMatch(id -> id == null)) { + throw new IllegalArgumentException("At least one notice receiver is required"); + } + schedule.setReceiverIds(receiverIds); + List receivers = receiverIds.stream() + .map(id -> { + NoticeReceiver receiver = noticeConfigService.getReceiverById(id); + if (receiver == null) { + throw new IllegalArgumentException("Notice receiver not found: " + id); + } + return receiver; + }) + .toList(); + if (schedule.getTemplateId() == null) { + for (NoticeReceiver receiver : receivers) { + if (receiver.getType() != 0 + && noticeConfigService.getDefaultNoticeTemplateByType(receiver.getType()) == null) { + throw new IllegalArgumentException( + "Default notice template not found for receiver type: " + receiver.getType()); + } + } + nextTriggerAt(schedule.getCronExpression()); + return; + } + NoticeTemplate template = noticeConfigService.getNoticeTemplatesById(schedule.getTemplateId()) + .orElseThrow(() -> new IllegalArgumentException( + "Notice template not found: " + schedule.getTemplateId())); + if (receivers.stream().anyMatch(receiver -> !template.getType().equals(receiver.getType()))) { + throw new IllegalArgumentException("Notice template type must match every receiver type"); + } + nextTriggerAt(schedule.getCronExpression()); + } + + private long nextTriggerAt(String expression) { + // Cron input may contain repeated whitespace; split fields only to enforce the minute-level scheduling contract. + String[] fields = expression.trim().split("\\s+"); + if (fields.length != 6 || !"0".equals(fields[0])) { + throw new IllegalArgumentException( + "Agent schedule requires a six-field Spring cron expression with seconds set to 0"); + } + ZonedDateTime next = CronExpression.parse(expression) + .next(ZonedDateTime.now(ZoneId.systemDefault())); + if (next == null) { + throw new IllegalArgumentException("Agent schedule cron expression has no next execution time"); + } + return next.toInstant().toEpochMilli(); + } +} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutor.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutor.java deleted file mode 100644 index 4954d49740..0000000000 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutor.java +++ /dev/null @@ -1,89 +0,0 @@ -/* - * 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.gateway.schedule; - -import java.util.List; -import lombok.extern.slf4j.Slf4j; -import tools.jackson.core.type.TypeReference; -import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.InvokeCommand; -import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.ReplyMode; -import org.apache.hertzbeat.ai.gateway.application.GatewayCommandRouter; -import org.apache.hertzbeat.ai.gateway.contract.AgentResponseLanguage; -import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; -import org.apache.hertzbeat.ai.gateway.contract.UserInput; -import org.apache.hertzbeat.ai.gateway.identity.AgentActor; -import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; -import org.apache.hertzbeat.common.util.JsonUtil; -import org.springframework.scheduling.annotation.Scheduled; -import org.springframework.stereotype.Component; - -/** - * Routes due scheduled commands through the standard Gateway command path. - */ -@Slf4j -@Component -public class AgentScheduledCommandExecutor { - - private final AgentScheduledCommandService commandService; - private final GatewayCommandRouter commandRouter; - - public AgentScheduledCommandExecutor(AgentScheduledCommandService commandService, - GatewayCommandRouter commandRouter) { - this.commandService = commandService; - this.commandRouter = commandRouter; - } - - @Scheduled(fixedDelay = 60_000L) - public void executeDueCommands() { - for (AgentScheduledCommand scheduled : commandService.findDueCommands()) { - try { - commandRouter.handle(command(scheduled)); - } catch (RuntimeException exception) { - log.error("Scheduled Gateway command {} failed", scheduled.getId(), exception); - } finally { - commandService.completeExecution(scheduled); - } - } - } - - private InvokeCommand command(AgentScheduledCommand scheduled) { - List roles = JsonUtil.fromJson(scheduled.getActorRoles(), new TypeReference>() { }); - long now = System.currentTimeMillis(); - String commandId = "scheduled_" + scheduled.getId() + "_" + now; - return InvokeCommand.builder() - .envelope(GatewayEnvelope.builder() - .channelId(scheduled.getChannel()) - .receivedAt(now) - .preferredLanguage(AgentResponseLanguage.systemDefault()) - .actor(AgentActor.builder() - .type(scheduled.getActorType()) - .id(scheduled.getActorId()) - .roles(roles) - .build()) - .build()) - .replyMode(ReplyMode.FINAL_ONLY) - .commandId(commandId) - .entryType(AgentRuntimeEntryType.SCHEDULE_TRIGGER) - .userInput(UserInput.builder() - .messageId(commandId) - .conversationId(scheduled.getConversationId()) - .message(UserInput.Message.builder().text(scheduled.getMessage()).build()) - .build()) - .build(); - } -} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandService.java deleted file mode 100644 index 11e5777963..0000000000 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandService.java +++ /dev/null @@ -1,81 +0,0 @@ -/* - * 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.gateway.schedule; - -import java.time.LocalDateTime; -import java.util.List; -import org.springframework.scheduling.support.CronExpression; -import org.springframework.stereotype.Service; -import org.springframework.transaction.annotation.Transactional; - -/** - * Manages recurring Gateway commands. - */ -@Service -public class AgentScheduledCommandService { - - private final AgentScheduledCommandDao commandDao; - - public AgentScheduledCommandService(AgentScheduledCommandDao commandDao) { - this.commandDao = commandDao; - } - - @Transactional - public AgentScheduledCommand create(AgentScheduledCommand command) { - command.setNextRunTime(command.isEnabled() ? next(command.getCronExpression()) : null); - return commandDao.save(command); - } - - public List findSessionCommands(Long sessionId) { - return commandDao.findBySessionIdOrderByIdAsc(sessionId); - } - - public List findDueCommands() { - return commandDao.findByEnabledTrueAndNextRunTimeLessThanEqual(LocalDateTime.now()); - } - - public AgentScheduledCommand getOwned(Long commandId, Long sessionId) { - return commandDao.findById(commandId) - .filter(command -> sessionId.equals(command.getSessionId())) - .orElseThrow(() -> new IllegalArgumentException("Scheduled command not found in the current Agent session")); - } - - @Transactional - public void delete(Long commandId, Long sessionId) { - commandDao.delete(getOwned(commandId, sessionId)); - } - - @Transactional - public AgentScheduledCommand toggle(Long commandId, Long sessionId, boolean enabled) { - AgentScheduledCommand command = getOwned(commandId, sessionId); - command.setEnabled(enabled); - command.setNextRunTime(enabled ? next(command.getCronExpression()) : null); - return commandDao.save(command); - } - - @Transactional - public void completeExecution(AgentScheduledCommand command) { - command.setLastRunTime(LocalDateTime.now()); - command.setNextRunTime(command.isEnabled() ? next(command.getCronExpression()) : null); - commandDao.save(command); - } - - private LocalDateTime next(String expression) { - return CronExpression.parse(expression).next(LocalDateTime.now()); - } -} diff --git a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/tool/schedule/AgentScheduleToolService.java b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/tool/schedule/AgentScheduleToolService.java index 405aa22581..cc54454253 100644 --- a/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/tool/schedule/AgentScheduleToolService.java +++ b/hertzbeat-ai-gateway/src/main/java/org/apache/hertzbeat/ai/gateway/tool/schedule/AgentScheduleToolService.java @@ -18,17 +18,18 @@ package org.apache.hertzbeat.ai.gateway.tool.schedule; import java.util.List; -import org.apache.hertzbeat.ai.gateway.conversation.AgentSessionService; -import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; -import org.apache.hertzbeat.ai.gateway.schedule.AgentScheduledCommand; -import org.apache.hertzbeat.ai.gateway.schedule.AgentScheduledCommandService; +import org.apache.hertzbeat.ai.gateway.schedule.AgentSchedule; +import org.apache.hertzbeat.ai.gateway.schedule.AgentScheduleExecutor; +import org.apache.hertzbeat.ai.gateway.schedule.AgentScheduleService; import org.apache.hertzbeat.ai.gateway.tool.core.AgentToolContextSupport; import org.apache.hertzbeat.ai.gateway.tool.core.AgentToolExposure; import org.apache.hertzbeat.ai.gateway.tool.core.AgentToolPolicy; import org.apache.hertzbeat.ai.gateway.tool.core.AgentToolRisk; -import org.apache.hertzbeat.common.entity.agent.AgentSession; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.springframework.data.domain.PageRequest; import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.annotation.ToolParam; +import org.springframework.context.annotation.Lazy; import org.springframework.stereotype.Service; /** @@ -37,43 +38,65 @@ import org.springframework.stereotype.Service; @Service public class AgentScheduleToolService { - private final AgentScheduledCommandService commandService; - private final AgentSessionService sessionService; + private final AgentScheduleService scheduleService; + private final AgentScheduleExecutor scheduleExecutor; - public AgentScheduleToolService(AgentScheduledCommandService commandService, - AgentSessionService sessionService) { - this.commandService = commandService; - this.sessionService = sessionService; + public AgentScheduleToolService(AgentScheduleService scheduleService, + @Lazy AgentScheduleExecutor scheduleExecutor) { + this.scheduleService = scheduleService; + this.scheduleExecutor = scheduleExecutor; } - @Tool(name = "schedule.create", description = "Schedule a recurring user request in the current Agent session.") + @Tool(name = "schedule.create", description = "Create a system-level recurring Agent inspection.") @AgentToolPolicy(risk = AgentToolRisk.CHANGE, exposure = AgentToolExposure.MODEL_ON_DEMAND) - public AgentScheduledCommand createSchedule( - @ToolParam(description = "User request to run at each scheduled time.") String message, - @ToolParam(description = "Six-field Spring cron expression.") String cronExpression, + public AgentSchedule createSchedule( + @ToolParam(description = "Schedule name.") String name, + @ToolParam(description = "Inspection instruction to run at each scheduled time.") String instruction, + @ToolParam(description = "Six-field Spring cron expression with seconds set to 0.") + String cronExpression, + @ToolParam(description = "Existing HertzBeat notice receiver IDs.") List receiverIds, + @ToolParam(required = false, description = "Existing notice template ID; omit to use channel defaults.") + Long templateId, @ToolParam(required = false, description = "Whether the schedule starts enabled.") Boolean enabled) { var request = AgentToolContextSupport.invocation().getRequest(); - AgentSession session = sessionService.findSession(request.getSessionUid()) - .orElseThrow(() -> new IllegalStateException("Current Agent session was not found")); - return commandService.create(AgentScheduledCommand.builder() - .sessionId(session.getId()) - .channel(session.getChannel()) - .conversationId(session.getConversationId()) - .actorType(request.getActor().getType()) - .actorId(request.getActor().getId()) - .actorRoles(ActorSupport.rolesJson(request.getActor())) - .message(message) + return scheduleService.create(AgentSchedule.builder() + .name(name) + .instruction(instruction) .cronExpression(cronExpression) .enabled(enabled == null || enabled) + .receiverIds(receiverIds) + .templateId(templateId) + .createdFromSessionUid(request.getSessionUid()) .build()); } - @Tool(name = "schedule.list", description = "List recurring Gateway commands for the current session.") + @Tool(name = "schedule.list", description = "List system-level Agent schedules.") @AgentToolPolicy( exposure = AgentToolExposure.MODEL_ON_DEMAND) - public List listSchedules() { - return commandService.findSessionCommands(AgentToolContextSupport.invocation().getRequest().getRunSessionId()); + public List listSchedules() { + return scheduleService.list(PageRequest.of(0, 100)).getContent(); + } + + @Tool(name = "schedule.update", description = "Update an existing system-level Agent inspection.") + @AgentToolPolicy(risk = AgentToolRisk.CHANGE, + exposure = AgentToolExposure.MODEL_ON_DEMAND) + public AgentSchedule updateSchedule( + @ToolParam(description = "Schedule id.") Long scheduleId, + @ToolParam(description = "Schedule name.") String name, + @ToolParam(description = "Inspection instruction to run at each scheduled time.") String instruction, + @ToolParam(description = "Six-field Spring cron expression with seconds set to 0.") + String cronExpression, + @ToolParam(description = "Existing HertzBeat notice receiver IDs.") List receiverIds, + @ToolParam(required = false, description = "Existing notice template ID; omit to use channel defaults.") + Long templateId) { + return scheduleService.update(scheduleId, AgentSchedule.builder() + .name(name) + .instruction(instruction) + .cronExpression(cronExpression) + .receiverIds(receiverIds) + .templateId(templateId) + .build()); } @Tool(name = "schedule.delete", description = "Delete a recurring diagnostic schedule.") @@ -86,17 +109,24 @@ public class AgentScheduleToolService { if (reason == null || reason.isBlank()) { throw new IllegalArgumentException("reason is required for schedule.delete"); } - commandService.delete(scheduleId, AgentToolContextSupport.invocation().getRequest().getRunSessionId()); + scheduleService.delete(scheduleId); return "Schedule deleted: " + scheduleId; } @Tool(name = "schedule.toggle", description = "Enable or disable a recurring diagnostic schedule.") @AgentToolPolicy(risk = AgentToolRisk.CHANGE, exposure = AgentToolExposure.MODEL_ON_DEMAND) - public AgentScheduledCommand toggleSchedule( + public AgentSchedule toggleSchedule( @ToolParam(description = "Schedule id.") Long scheduleId, @ToolParam(description = "Whether the schedule should be enabled.") boolean enabled) { - return commandService.toggle(scheduleId, - AgentToolContextSupport.invocation().getRequest().getRunSessionId(), enabled); + return scheduleService.toggle(scheduleId, enabled); } + + @Tool(name = "schedule.run_now", description = "Run a system-level Agent schedule immediately.") + @AgentToolPolicy(risk = AgentToolRisk.CHANGE, + exposure = AgentToolExposure.MODEL_ON_DEMAND) + public AgentRun runNow(@ToolParam(description = "Schedule id.") Long scheduleId) { + return scheduleExecutor.executeNow(scheduleId); + } + } diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandlerTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandlerTest.java index c94b732ec1..c6b3b5b1e4 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandlerTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/alert/AgentAlertAnalysisEventHandlerTest.java @@ -29,6 +29,7 @@ import java.util.Map; import org.apache.hertzbeat.alert.service.AlertAnalysisPolicyService; import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.InvokeCommand; import org.apache.hertzbeat.ai.gateway.application.GatewayCommandRouter; +import org.apache.hertzbeat.ai.gateway.channel.core.ChannelId; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; import org.apache.hertzbeat.common.entity.alerter.AlertAnalysisPolicy; @@ -86,6 +87,7 @@ class AgentAlertAnalysisEventHandlerTest { ArgumentCaptor command = ArgumentCaptor.forClass(InvokeCommand.class); verify(commandRouter, org.mockito.Mockito.timeout(2000)).handle(command.capture()); + assertEquals(ChannelId.SYSTEM.id(), command.getValue().envelope().getChannelId()); assertEquals(AgentRuntimeEntryType.ALERT_TRIGGER, command.getValue().entryType()); assertEquals(List.of(ActorSupport.ROLE_ALERT_ANALYSIS), command.getValue().envelope().getActor().getRoles()); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandServiceTest.java index f3476cfa02..08e698be9b 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandServiceTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/AgentCommandServiceTest.java @@ -76,8 +76,8 @@ class AgentCommandServiceTest { void setUp() { session = AgentSession.builder().id(1L).sessionUid("session-1").build(); run = AgentRun.builder().id(2L).runUid("run-1").sessionId(1L).build(); - when(sessionService.findOrCreateSession(any(), any())).thenReturn(session); - when(runService.createOrResumeRun(any(), any())).thenReturn(run); + when(sessionService.findOrCreateSession(any(), any(), any())).thenReturn(session); + when(runService.createOrResumeRun(any(), any(), any())).thenReturn(run); when(transcriptRecorder.chatHistory(session.getId())).thenReturn(List.of()); when(runService.markRunning(run)).thenReturn(run); } diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommandRouterTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommandRouterTest.java index 4562ddb1fa..ad2c35c3b0 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommandRouterTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayCommandRouterTest.java @@ -37,6 +37,7 @@ import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; import org.apache.hertzbeat.ai.gateway.contract.UserInput; import org.apache.hertzbeat.ai.gateway.contract.UserInput.Message; import org.apache.hertzbeat.ai.gateway.identity.AgentActor; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.common.entity.dto.ModelProviderConfig; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; @@ -81,7 +82,8 @@ class GatewayCommandRouterTest { void queryCommandShouldRouteOnlyToQueryService() { GatewaySingleResponse expected = response("session"); GetSessionCommand command = new GetSessionCommand( - envelope(), ReplyMode.FINAL_ONLY, "get-session", "ags-1"); + envelope(), ReplyMode.FINAL_ONLY, "get-session", + AgentRuntimeEntryType.USER_INPUT, "ags-1"); when(queryService.getSession(command)).thenReturn(expected); assertSame(expected, router().handle(command)); @@ -95,7 +97,8 @@ class GatewayCommandRouterTest { void sessionListCommandShouldRouteOnlyToQueryService() { GatewaySingleResponse expected = response("sessions"); ListSessionsCommand command = new ListSessionsCommand( - envelope(), ReplyMode.FINAL_ONLY, "list-sessions", null, 0, 50); + envelope(), ReplyMode.FINAL_ONLY, "list-sessions", + AgentRuntimeEntryType.USER_INPUT, null, 0, 50); when(queryService.listSessions(command)).thenReturn(expected); assertSame(expected, router().handle(command)); @@ -109,7 +112,8 @@ class GatewayCommandRouterTest { void transcriptCommandShouldRouteOnlyToQueryService() { GatewaySingleResponse expected = response("session-transcript"); GetSessionTranscriptCommand command = new GetSessionTranscriptCommand( - envelope(), ReplyMode.FINAL_ONLY, "get-transcript", "ags-1", 0, 50); + envelope(), ReplyMode.FINAL_ONLY, "get-transcript", + AgentRuntimeEntryType.USER_INPUT, "ags-1", 0, 50); when(queryService.getSessionTranscript(command)).thenReturn(expected); assertSame(expected, router().handle(command)); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryServiceTest.java index 45a6e66ba5..eb85c6015f 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryServiceTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/application/GatewayQueryServiceTest.java @@ -33,6 +33,7 @@ import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.ReplyMode; import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; import org.apache.hertzbeat.ai.gateway.conversation.AgentSessionService; import org.apache.hertzbeat.ai.gateway.identity.AgentActor; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.common.entity.agent.AgentSession; import org.apache.hertzbeat.common.entity.agent.AgentTranscriptEntry; import org.junit.jupiter.api.Test; @@ -57,41 +58,49 @@ class GatewayQueryServiceTest { GatewayEnvelope envelope = envelope("bob"); PageRequest pageRequest = PageRequest.of(0, 50); Page sessions = Page.empty(pageRequest); - when(sessionService.findSessions(envelope, null, pageRequest)).thenReturn(sessions); + when(sessionService.findSessions( + envelope, AgentRuntimeEntryType.USER_INPUT, null, pageRequest)).thenReturn(sessions); ListSessionsCommand command = new ListSessionsCommand( - envelope, ReplyMode.FINAL_ONLY, "list-sessions", null, 0, 50); + envelope, ReplyMode.FINAL_ONLY, "list-sessions", + AgentRuntimeEntryType.USER_INPUT, null, 0, 50); GatewayResponse.GatewaySingleResponse response = service().listSessions(command); assertSame(sessions, response.body()); - verify(sessionService).findSessions(envelope, null, pageRequest); + verify(sessionService).findSessions( + envelope, AgentRuntimeEntryType.USER_INPUT, null, pageRequest); } @Test void listSessionsShouldForwardAlertAnalysisEnvelopeAndTitle() { GatewayEnvelope envelope = GatewayEnvelope.builder() - .channelId("alert") + .channelId("system") .receivedAt(100L) .actor(AgentActor.alertAnalysisActor()) .build(); PageRequest pageRequest = PageRequest.of(0, 50); Page sessions = Page.empty(pageRequest); - when(sessionService.findSessions(envelope, "database", pageRequest)).thenReturn(sessions); + when(sessionService.findSessions( + envelope, AgentRuntimeEntryType.ALERT_TRIGGER, "database", pageRequest)).thenReturn(sessions); ListSessionsCommand command = new ListSessionsCommand( - envelope, ReplyMode.FINAL_ONLY, "list-alert-sessions", "database", 0, 50); + envelope, ReplyMode.FINAL_ONLY, "list-alert-sessions", + AgentRuntimeEntryType.ALERT_TRIGGER, "database", 0, 50); GatewayResponse.GatewaySingleResponse response = service().listSessions(command); assertSame(sessions, response.body()); - verify(sessionService).findSessions(envelope, "database", pageRequest); + verify(sessionService).findSessions( + envelope, AgentRuntimeEntryType.ALERT_TRIGGER, "database", pageRequest); } @Test void getSessionShouldHideAnotherWebUiActorsSession() { GatewayEnvelope envelope = envelope("bob"); - when(sessionService.findOwnedSession("ags-alice", envelope)).thenReturn(Optional.empty()); + when(sessionService.findOwnedSession( + "ags-alice", envelope, AgentRuntimeEntryType.USER_INPUT)).thenReturn(Optional.empty()); GetSessionCommand command = new GetSessionCommand( - envelope, ReplyMode.FINAL_ONLY, "get-session", "ags-alice"); + envelope, ReplyMode.FINAL_ONLY, "get-session", + AgentRuntimeEntryType.USER_INPUT, "ags-alice"); GatewayResponse.GatewaySingleResponse response = service().getSession(command); @@ -103,9 +112,11 @@ class GatewayQueryServiceTest { @Test void getTranscriptShouldNotLoadAnotherWebUiActorsEntries() { GatewayEnvelope envelope = envelope("bob"); - when(sessionService.findOwnedSession("ags-alice", envelope)).thenReturn(Optional.empty()); + when(sessionService.findOwnedSession( + "ags-alice", envelope, AgentRuntimeEntryType.USER_INPUT)).thenReturn(Optional.empty()); GetSessionTranscriptCommand command = new GetSessionTranscriptCommand( - envelope, ReplyMode.FINAL_ONLY, "get-transcript", "ags-alice", 0, 50); + envelope, ReplyMode.FINAL_ONLY, "get-transcript", + AgentRuntimeEntryType.USER_INPUT, "ags-alice", 0, 50); GatewayResponse.GatewaySingleResponse response = service().getSessionTranscript(command); @@ -122,10 +133,12 @@ class GatewayQueryServiceTest { PageRequest pageRequest = PageRequest.of(0, 50); Page transcript = new PageImpl<>( List.of(AgentTranscriptEntry.builder().sessionId(2L).build()), pageRequest, 1); - when(sessionService.findOwnedSession("ags-bob", envelope)).thenReturn(Optional.of(session)); + when(sessionService.findOwnedSession( + "ags-bob", envelope, AgentRuntimeEntryType.USER_INPUT)).thenReturn(Optional.of(session)); when(sessionService.findTranscriptEntries(2L, pageRequest)).thenReturn(transcript); GetSessionTranscriptCommand command = new GetSessionTranscriptCommand( - envelope, ReplyMode.FINAL_ONLY, "get-transcript", "ags-bob", 0, 50); + envelope, ReplyMode.FINAL_ONLY, "get-transcript", + AgentRuntimeEntryType.USER_INPUT, "ags-bob", 0, 50); GatewayResponse.GatewaySingleResponse response = service().getSessionTranscript(command); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryControllerTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryControllerTest.java index 2222aa1f9c..31aa568d28 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryControllerTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/channel/webui/QueryControllerTest.java @@ -37,6 +37,7 @@ import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.GetSessionTran import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.ListSessionsCommand; import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.Meta; import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.GatewaySingleResponse; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.common.entity.agent.AgentSession; import org.apache.hertzbeat.common.entity.agent.AgentTranscriptEntry; import org.apache.hertzbeat.common.entity.dto.Message; @@ -87,6 +88,7 @@ class QueryControllerTest { GetSessionCommand command = (GetSessionCommand) commandCaptor.getValue(); assertEquals(ChannelId.WEB_UI.id(), command.envelope().getChannelId()); assertEquals("trusted-user", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.USER_INPUT, command.originEntryType()); } @Test @@ -103,6 +105,7 @@ class QueryControllerTest { assertInstanceOf(ListSessionsCommand.class, commandCaptor.getValue()); ListSessionsCommand command = (ListSessionsCommand) commandCaptor.getValue(); assertEquals("trusted-user", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.USER_INPUT, command.originEntryType()); assertEquals(0, command.pageIndex()); assertEquals(50, command.pageSize()); } @@ -120,9 +123,10 @@ class QueryControllerTest { assertSame(page, response.getBody().getData()); assertInstanceOf(ListSessionsCommand.class, commandCaptor.getValue()); ListSessionsCommand command = (ListSessionsCommand) commandCaptor.getValue(); - assertEquals(ChannelId.ALERT.id(), command.envelope().getChannelId()); + assertEquals(ChannelId.SYSTEM.id(), command.envelope().getChannelId()); assertEquals("system", command.envelope().getActor().getType()); assertEquals("alert-analysis", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.ALERT_TRIGGER, command.originEntryType()); assertEquals("database", command.title()); } @@ -136,8 +140,9 @@ class QueryControllerTest { assertInstanceOf(GetSessionCommand.class, commandCaptor.getValue()); GetSessionCommand command = (GetSessionCommand) commandCaptor.getValue(); - assertEquals(ChannelId.ALERT.id(), command.envelope().getChannelId()); + assertEquals(ChannelId.SYSTEM.id(), command.envelope().getChannelId()); assertEquals("alert-analysis", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.ALERT_TRIGGER, command.originEntryType()); assertEquals("ags-1", command.sessionUid()); } @@ -161,6 +166,7 @@ class QueryControllerTest { assertInstanceOf(GetSessionTranscriptCommand.class, commandCaptor.getValue()); GetSessionTranscriptCommand command = (GetSessionTranscriptCommand) commandCaptor.getValue(); assertEquals("trusted-user", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.USER_INPUT, command.originEntryType()); assertEquals("ags-1", command.sessionUid()); assertEquals(0, command.pageIndex()); assertEquals(50, command.pageSize()); @@ -181,8 +187,9 @@ class QueryControllerTest { assertSame(transcript, response.getBody().getData()); assertInstanceOf(GetSessionTranscriptCommand.class, commandCaptor.getValue()); GetSessionTranscriptCommand command = (GetSessionTranscriptCommand) commandCaptor.getValue(); - assertEquals(ChannelId.ALERT.id(), command.envelope().getChannelId()); + assertEquals(ChannelId.SYSTEM.id(), command.envelope().getChannelId()); assertEquals("alert-analysis", command.envelope().getActor().getId()); + assertEquals(AgentRuntimeEntryType.ALERT_TRIGGER, command.originEntryType()); assertEquals("ags-1", command.sessionUid()); } diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunServiceTest.java index 73021574da..c5f410783c 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunServiceTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentRunServiceTest.java @@ -40,6 +40,7 @@ import org.apache.hertzbeat.ai.gateway.contract.AgentTargetRef; import org.apache.hertzbeat.ai.gateway.contract.AgentTopologyRef; import org.apache.hertzbeat.ai.gateway.contract.UserInput; import org.apache.hertzbeat.ai.gateway.contract.UserInput.Message; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.common.entity.agent.AgentRun; import org.apache.hertzbeat.common.entity.agent.AgentSession; import org.junit.jupiter.api.Test; @@ -74,7 +75,8 @@ class AgentRunServiceTest { AgentRun persisted = AgentRun.builder().id(2L).runUid("run_saved").build(); when(runDao.saveAndFlush(any(AgentRun.class))).thenReturn(persisted); - AgentRun result = service.createOrResumeRun(session, userInput); + AgentRun result = service.createOrResumeRun( + session, userInput, AgentRuntimeEntryType.USER_INPUT); ArgumentCaptor captor = ArgumentCaptor.forClass(AgentRun.class); verify(runDao).saveAndFlush(captor.capture()); @@ -88,6 +90,7 @@ class AgentRunServiceTest { assertTrue(saved.getRunUid().startsWith("run_")); assertSame(persisted, result); assertEquals(AgentRunStatus.CREATED.name(), saved.getStatus()); + assertEquals(AgentRuntimeEntryType.USER_INPUT.name(), saved.getEntryType()); assertEquals("msg_1", saved.getMessageId()); } @@ -103,7 +106,8 @@ class AgentRunServiceTest { AgentRun existed = AgentRun.builder().id(2L).runUid("run_1").messageId("msg_1").build(); when(runDao.findBySessionIdAndMessageId(1L, "msg_1")).thenReturn(Optional.of(existed)); - AgentRun result = service.createOrResumeRun(session, userInput); + AgentRun result = service.createOrResumeRun( + session, userInput, AgentRuntimeEntryType.ALERT_TRIGGER); assertSame(existed, result); verify(runDao, never()).saveAndFlush(any()); @@ -138,7 +142,7 @@ class AgentRunServiceTest { when(runDao.findBySessionIdAndMessageId(1L, "msg_context")).thenReturn(Optional.empty()); when(runDao.saveAndFlush(any(AgentRun.class))).thenAnswer(invocation -> invocation.getArgument(0)); - AgentRun saved = service.createOrResumeRun(session, userInput); + AgentRun saved = service.createOrResumeRun(session, userInput, AgentRuntimeEntryType.USER_INPUT); AgentTargetRef persisted = service.targetFromRun(saved); assertEquals(300L, persisted.getEntityId()); @@ -161,7 +165,8 @@ class AgentRunServiceTest { when(runDao.findBySessionIdAndMessageId(1L, "msg_1")).thenReturn(Optional.empty(), Optional.of(existed)); when(runDao.saveAndFlush(any(AgentRun.class))).thenThrow(new DataIntegrityViolationException("duplicate")); - AgentRun result = service.createOrResumeRun(session, userInput); + AgentRun result = service.createOrResumeRun( + session, userInput, AgentRuntimeEntryType.SCHEDULE_TRIGGER); assertSame(existed, result); verify(entityManager).clear(); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionServiceTest.java index 0ae6a8895f..c9b54b5fd9 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionServiceTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/conversation/AgentSessionServiceTest.java @@ -42,6 +42,7 @@ import org.apache.hertzbeat.ai.gateway.contract.UserInput.Message; import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; import org.apache.hertzbeat.ai.gateway.contract.UserInput; import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeHistoryWindow; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; import org.apache.hertzbeat.ai.gateway.runtime.TranscriptContent; import org.apache.hertzbeat.ai.gateway.runtime.TranscriptMessage; import org.apache.hertzbeat.common.entity.agent.AgentSession; @@ -81,13 +82,14 @@ class AgentSessionServiceTest { sessionDao, transcriptEntryDao, sessionKeyBuilder, entityManager); GatewayEnvelope envelope = envelope("bob"); AgentSession session = AgentSession.builder().id(1L).sessionUid("ags-1").build(); - when(sessionDao.findBySessionUidAndChannelAndActorTypeAndActorId( - "ags-1", "web-ui", "user", "bob")).thenReturn(Optional.of(session)); + when(sessionDao.findBySessionUidAndChannelAndActorTypeAndActorIdAndOriginEntryType( + "ags-1", "web-ui", "user", "bob", "USER_INPUT")).thenReturn(Optional.of(session)); - assertSame(session, service.findOwnedSession(" ags-1 ", envelope).orElseThrow()); + assertSame(session, service.findOwnedSession( + " ags-1 ", envelope, AgentRuntimeEntryType.USER_INPUT).orElseThrow()); - verify(sessionDao).findBySessionUidAndChannelAndActorTypeAndActorId( - "ags-1", "web-ui", "user", "bob"); + verify(sessionDao).findBySessionUidAndChannelAndActorTypeAndActorIdAndOriginEntryType( + "ags-1", "web-ui", "user", "bob", "USER_INPUT"); } @Test @@ -96,13 +98,14 @@ class AgentSessionServiceTest { sessionDao, transcriptEntryDao, sessionKeyBuilder, entityManager); GatewayEnvelope envelope = envelope("bob"); AgentSession session = AgentSession.builder().id(42L).sessionUid("ags-42").build(); - when(sessionDao.findByIdAndChannelAndActorTypeAndActorId( - 42L, "web-ui", "user", "bob")).thenReturn(Optional.of(session)); + when(sessionDao.findByIdAndChannelAndActorTypeAndActorIdAndOriginEntryType( + 42L, "web-ui", "user", "bob", "USER_INPUT")).thenReturn(Optional.of(session)); - assertSame(session, service.findOwnedSession("42", envelope).orElseThrow()); + assertSame(session, service.findOwnedSession( + "42", envelope, AgentRuntimeEntryType.USER_INPUT).orElseThrow()); - verify(sessionDao).findByIdAndChannelAndActorTypeAndActorId( - 42L, "web-ui", "user", "bob"); + verify(sessionDao).findByIdAndChannelAndActorTypeAndActorIdAndOriginEntryType( + 42L, "web-ui", "user", "bob", "USER_INPUT"); } @Test @@ -110,19 +113,22 @@ class AgentSessionServiceTest { AgentSessionService service = new AgentSessionService( sessionDao, transcriptEntryDao, sessionKeyBuilder, entityManager); GatewayEnvelope envelope = GatewayEnvelope.builder() - .channelId("alert") + .channelId("system") .receivedAt(100L) .actor(AgentActor.alertAnalysisActor()) .build(); PageRequest pageRequest = PageRequest.of(0, 8); Page page = Page.empty(pageRequest); - when(sessionDao.findByChannelAndActorTypeAndActorIdAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( - "alert", "system", "alert-analysis", "database", pageRequest)).thenReturn(page); + when(sessionDao + .findByChannelAndActorTypeAndActorIdAndOriginEntryTypeAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( + "system", "system", "alert-analysis", "ALERT_TRIGGER", "database", pageRequest)).thenReturn(page); - assertSame(page, service.findSessions(envelope, "database", pageRequest)); + assertSame(page, service.findSessions( + envelope, AgentRuntimeEntryType.ALERT_TRIGGER, "database", pageRequest)); - verify(sessionDao).findByChannelAndActorTypeAndActorIdAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( - "alert", "system", "alert-analysis", "database", pageRequest); + verify(sessionDao) + .findByChannelAndActorTypeAndActorIdAndOriginEntryTypeAndTitleContainingIgnoreCaseOrderByGmtUpdateDesc( + "system", "system", "alert-analysis", "ALERT_TRIGGER", "database", pageRequest); } @Test @@ -150,7 +156,8 @@ class AgentSessionServiceTest { when(sessionKeyBuilder.build(envelope, "chat-1")).thenReturn("key-1"); when(sessionDao.findBySessionKey("key-1")).thenReturn(Optional.of(existing)); - AgentSession session = service.findOrCreateSession(envelope, userInput); + AgentSession session = service.findOrCreateSession( + envelope, userInput, AgentRuntimeEntryType.USER_INPUT); assertSame(existing, session); assertEquals("[\"old-role\"]", existing.getActorRoles()); @@ -177,10 +184,12 @@ class AgentSessionServiceTest { when(sessionDao.saveAndFlush(any(AgentSession.class))) .thenAnswer(invocation -> invocation.getArgument(0)); - AgentSession session = service.findOrCreateSession(envelope, userInput); + AgentSession session = service.findOrCreateSession( + envelope, userInput, AgentRuntimeEntryType.USER_INPUT); assertEquals(AgentSessionStatus.ACTIVE, session.getStatus()); assertEquals("ACTIVE", session.getStatus().name()); + assertEquals(AgentRuntimeEntryType.USER_INPUT.name(), session.getOriginEntryType()); } @Test @@ -254,6 +263,26 @@ class AgentSessionServiceTest { assertEquals(4L, result.getContent().get(0).getSessionSequence()); } + @Test + void findConversationTranscriptEntriesShouldQueryLatestUserAndAssistantMessages() { + AgentSessionService service = new AgentSessionService( + sessionDao, transcriptEntryDao, sessionKeyBuilder, entityManager); + PageRequest pageable = PageRequest.of(0, 20); + AgentTranscriptEntry entry = transcriptEntry(8L, TranscriptMessage.assistantText("latest", null)); + Page page = new PageImpl<>(List.of(entry), pageable, 21); + List roles = List.of( + TranscriptMessage.TranscriptRole.USER.wireValue(), + TranscriptMessage.TranscriptRole.ASSISTANT.wireValue()); + when(transcriptEntryDao.findBySessionIdAndMessageRoleInOrderBySessionSequenceDesc( + eq(1L), eq(roles), eq(pageable))).thenReturn(page); + + Page result = service.findConversationTranscriptEntries(1L, pageable); + + assertSame(page, result); + assertEquals(8L, result.getContent().get(0).getSessionSequence()); + assertEquals(2, result.getTotalPages()); + } + @Test void persistCompactionCheckpointShouldAppendDerivedSummary() { AgentSessionService service = new AgentSessionService( diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilderTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilderTest.java index f56472dcb7..dfe2ed9a3a 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilderTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeContextBuilderTest.java @@ -28,6 +28,7 @@ import java.time.Instant; import java.time.ZoneOffset; import java.util.List; import java.util.Map; +import java.util.TimeZone; import org.apache.hertzbeat.ai.gateway.contract.AgentSignalRef; import org.apache.hertzbeat.ai.gateway.contract.AgentTargetRef; import org.apache.hertzbeat.ai.gateway.contract.AgentTopologyRef; @@ -233,6 +234,29 @@ class AgentRuntimeContextBuilderTest { assertEquals(3, restored.getTopology().getDepth()); } + @Test + void shouldUseHertzBeatSystemTimezone() { + TimeZone original = TimeZone.getDefault(); + try { + TimeZone.setDefault(TimeZone.getTimeZone("Asia/Shanghai")); + AgentRuntimeRequest request = AgentRuntimeRequest.builder() + .envelope(envelope()) + .session(session()) + .run(run()) + .entryType(AgentRuntimeEntryType.USER_INPUT) + .approvalHandling(AgentApprovalHandling.WAIT_FOR_DECISION) + .userInput(userInput("inspect")) + .build(); + + AgentRuntimeContext context = builder("trace").build(request, new AgentRuntimeProperties()); + + assertEquals("Asia/Shanghai", context.getTimezone()); + assertEquals("2026-04-19T08:00:00+08:00", context.getCurrentTimeIso()); + } finally { + TimeZone.setDefault(original); + } + } + private AgentRuntimeContextBuilder builder(String traceId) { return new AgentRuntimeContextBuilder(Clock.fixed(NOW, ZoneOffset.UTC), () -> traceId); } diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeLoopTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeLoopTest.java index 5fa1acc08a..4f25ea2db2 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeLoopTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/AgentRuntimeLoopTest.java @@ -345,7 +345,7 @@ class AgentRuntimeLoopTest { assertEquals(modelArguments(), catalog.lastRequest.getArguments()); AgentRuntimeModelRequest secondRequest = modelClient.requests.get(1); String runtimeContext = runtimeContext(secondRequest.getPrompt()); - assertTrue(runtimeContext.contains("Current time: 2026-04-19T00:00:00Z")); + assertTrue(runtimeContext.contains("Current time: " + context.getCurrentTimeIso())); assertFalse(runtimeContext.contains("Run: uid=run-context")); assertFalse(runtimeContext.contains("run-context")); assertFalse(runtimeContext.contains("session-context")); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/RuntimePromptBuilderTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/RuntimePromptBuilderTest.java index a11597cce9..1215e2da57 100644 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/RuntimePromptBuilderTest.java +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/runtime/RuntimePromptBuilderTest.java @@ -132,8 +132,8 @@ class RuntimePromptBuilderTest { assertFalse(runtimeContext.contains("Trusted")); assertFalse(runtimeContext.contains("")); assertTrue(runtimeContext.contains("### Time")); - assertTrue(runtimeContext.contains("Current time: 1970-01-01T00:00:00Z")); - assertTrue(runtimeContext.contains("Timezone: Z")); + assertTrue(runtimeContext.contains("Current time: " + context.getCurrentTimeIso())); + assertTrue(runtimeContext.contains("Timezone: " + context.getTimezone())); assertFalse(runtimeContext.contains("Deadline epoch millis")); assertTrue(runtimeContext.contains("### Channel")); assertTrue(runtimeContext.contains("Channel: web-ui")); diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutorTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutorTest.java new file mode 100644 index 0000000000..93d08fc939 --- /dev/null +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleExecutorTest.java @@ -0,0 +1,133 @@ +/* + * 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.gateway.schedule; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.timeout; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.InvokeCommand; +import org.apache.hertzbeat.ai.gateway.application.GatewayCommandRouter; +import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.GatewaySingleResponse; +import org.apache.hertzbeat.ai.gateway.application.GatewayResponse.Meta; +import org.apache.hertzbeat.ai.gateway.channel.core.ChannelId; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunService; +import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Tests for {@link AgentScheduleExecutor}. + */ +@ExtendWith(MockitoExtension.class) +class AgentScheduleExecutorTest { + + @Mock + private AgentScheduleService scheduleService; + @Mock + private AgentRunService runService; + @Mock + private GatewayCommandRouter commandRouter; + @Mock + private AgentScheduleNoticeService noticeService; + + private AgentScheduleExecutor executor; + + @AfterEach + void tearDown() { + if (executor != null) { + executor.close(); + } + } + + @Test + void shouldRunWithFixedSystemIdentityAndSession() throws InterruptedException { + AgentSchedule schedule = AgentSchedule.builder() + .id(7L) + .sessionId(21L) + .instruction("Inspect all unhealthy monitors") + .receiverIds(List.of(10L)) + .build(); + AgentRun run = AgentRun.builder() + .runUid("run_1") + .sessionId(21L) + .messageId("schedule:7:manual:1") + .status("CREATED") + .build(); + CountDownLatch handled = new CountDownLatch(1); + CountDownLatch release = new CountDownLatch(1); + when(scheduleService.claimManualRun(7L)).thenReturn(run); + when(scheduleService.get(7L)).thenReturn(schedule); + when(commandRouter.handle(any())).thenAnswer(invocation -> { + handled.countDown(); + release.await(5, TimeUnit.SECONDS); + return GatewaySingleResponse.builder() + .meta(Meta.builder().commandId(run.getMessageId()).terminal(true).message("completed").build()) + .body(Map.of("status", "SUCCEEDED", "message", "Healthy")) + .events(List.of()) + .build(); + }); + when(runService.findRun("run_1")).thenReturn(Optional.of(run)); + + executor = new AgentScheduleExecutor(scheduleService, runService, commandRouter, noticeService); + executor.executeNow(7L); + + assertTrue(handled.await(5, TimeUnit.SECONDS)); + assertThrows(IllegalStateException.class, () -> executor.executeNow(7L)); + release.countDown(); + verify(noticeService, timeout(5000)).send(schedule, run, true, "Healthy"); + ArgumentCaptor command = ArgumentCaptor.forClass(InvokeCommand.class); + verify(commandRouter).handle(command.capture()); + assertEquals(ChannelId.SYSTEM.id(), command.getValue().envelope().getChannelId()); + assertEquals(ActorSupport.TYPE_SYSTEM, command.getValue().envelope().getActor().getType()); + assertEquals(ActorSupport.ID_SCHEDULE, command.getValue().envelope().getActor().getId()); + assertEquals("schedule:7", command.getValue().userInput().getConversationId()); + assertEquals(AgentRuntimeEntryType.SCHEDULE_TRIGGER, command.getValue().entryType()); + } + + @Test + void shouldFailInterruptedRunOnStartup() { + AgentSchedule schedule = AgentSchedule.builder().id(7L).sessionId(21L).build(); + AgentRun running = AgentRun.builder().runUid("run_interrupted").status("RUNNING").build(); + AgentRun failed = AgentRun.builder().runUid("run_interrupted").status("FAILED").build(); + when(scheduleService.findInterrupted()).thenReturn(List.of(schedule)); + when(runService.findRunningRun(21L)).thenReturn(Optional.of(running)); + when(runService.markFailed(any(), any())).thenReturn(failed); + executor = new AgentScheduleExecutor(scheduleService, runService, commandRouter, noticeService); + + executor.failInterruptedRuns(); + + verify(runService).markFailed(running, "Agent schedule execution was interrupted by process restart"); + verify(noticeService).send(schedule, failed, false, + "Agent schedule execution was interrupted by process restart"); + } +} diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeServiceTest.java new file mode 100644 index 0000000000..62e600958f --- /dev/null +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleNoticeServiceTest.java @@ -0,0 +1,80 @@ +/* + * 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.gateway.schedule; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import org.apache.hertzbeat.alert.AlerterWorkerPool; +import org.apache.hertzbeat.alert.notice.AlertNoticeDispatch; +import org.apache.hertzbeat.alert.service.NoticeConfigService; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.apache.hertzbeat.common.entity.alerter.GroupAlert; +import org.apache.hertzbeat.common.entity.alerter.NoticeReceiver; +import org.apache.hertzbeat.common.entity.alerter.NoticeTemplate; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Tests for {@link AgentScheduleNoticeService}. + */ +@ExtendWith(MockitoExtension.class) +class AgentScheduleNoticeServiceTest { + + @Mock + private NoticeConfigService noticeConfigService; + @Mock + private AlertNoticeDispatch alertNoticeDispatch; + @Mock + private AlerterWorkerPool workerPool; + + @Test + void shouldUseExistingAlertNoticeTransportWithoutDispatchingAnAlert() { + NoticeReceiver receiver = NoticeReceiver.builder().id(10L).type((byte) 1).build(); + NoticeTemplate template = NoticeTemplate.builder().id(20L).type((byte) 1).build(); + when(noticeConfigService.getReceiverById(10L)).thenReturn(receiver); + when(noticeConfigService.getOneTemplateById(20L)).thenReturn(template); + doAnswer(invocation -> { + invocation.getArgument(1).run(); + return null; + }).when(workerPool).executeNotify(eq((byte) 1), any()); + AgentSchedule schedule = AgentSchedule.builder() + .id(7L) + .name("Daily inspection") + .receiverIds(List.of(10L)) + .templateId(20L) + .lastTriggerAt(100L) + .build(); + AgentRun run = AgentRun.builder().runUid("run_1").build(); + + new AgentScheduleNoticeService(noticeConfigService, alertNoticeDispatch, workerPool) + .send(schedule, run, true, "Everything is healthy"); + + ArgumentCaptor alert = ArgumentCaptor.forClass(GroupAlert.class); + verify(alertNoticeDispatch).sendNoticeMsg(eq(receiver), eq(template), alert.capture()); + assertEquals("resolved", alert.getValue().getStatus()); + assertEquals("Everything is healthy", alert.getValue().getAlerts().getFirst().getContent()); + } +} diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleServiceTest.java new file mode 100644 index 0000000000..5416b52a74 --- /dev/null +++ b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduleServiceTest.java @@ -0,0 +1,182 @@ +/* + * 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.gateway.schedule; + +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.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.List; +import java.util.Optional; +import org.apache.hertzbeat.alert.service.NoticeConfigService; +import org.apache.hertzbeat.ai.gateway.channel.core.ChannelId; +import org.apache.hertzbeat.ai.gateway.contract.GatewayEnvelope; +import org.apache.hertzbeat.ai.gateway.contract.UserInput; +import org.apache.hertzbeat.ai.gateway.conversation.AgentRunService; +import org.apache.hertzbeat.ai.gateway.conversation.AgentSessionService; +import org.apache.hertzbeat.ai.gateway.identity.ActorSupport; +import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; +import org.apache.hertzbeat.common.entity.agent.AgentRun; +import org.apache.hertzbeat.common.entity.agent.AgentSession; +import org.apache.hertzbeat.common.entity.alerter.NoticeReceiver; +import org.apache.hertzbeat.common.entity.alerter.NoticeTemplate; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +/** + * Tests for {@link AgentScheduleService}. + */ +@ExtendWith(MockitoExtension.class) +class AgentScheduleServiceTest { + + @Mock + private AgentScheduleDao scheduleDao; + @Mock + private AgentSessionService sessionService; + @Mock + private AgentRunService runService; + @Mock + private NoticeConfigService noticeConfigService; + + private AgentScheduleService service; + + @BeforeEach + void setUp() { + service = new AgentScheduleService(scheduleDao, sessionService, runService, noticeConfigService); + } + + @Test + void shouldCreateFixedSystemScheduleSession() { + AgentSchedule input = schedule(); + stubNoticeConfiguration(); + when(scheduleDao.saveAndFlush(input)).thenAnswer(invocation -> { + AgentSchedule saved = invocation.getArgument(0); + saved.setId(7L); + return saved; + }); + when(sessionService.findOrCreateSession(any(), any(), any())) + .thenReturn(AgentSession.builder().id(21L).build()); + when(scheduleDao.save(input)).thenReturn(input); + + AgentSchedule saved = service.create(input); + + assertEquals(21L, saved.getSessionId()); + assertNotNull(saved.getNextTriggerAt()); + ArgumentCaptor envelope = ArgumentCaptor.forClass(GatewayEnvelope.class); + ArgumentCaptor userInput = ArgumentCaptor.forClass(UserInput.class); + ArgumentCaptor entryType = ArgumentCaptor.forClass(AgentRuntimeEntryType.class); + verify(sessionService).findOrCreateSession( + envelope.capture(), userInput.capture(), entryType.capture()); + assertEquals(ChannelId.SYSTEM.id(), envelope.getValue().getChannelId()); + assertEquals(AgentRuntimeEntryType.SCHEDULE_TRIGGER, entryType.getValue()); + assertEquals(ActorSupport.TYPE_SYSTEM, envelope.getValue().getActor().getType()); + assertEquals(ActorSupport.ID_SCHEDULE, envelope.getValue().getActor().getId()); + assertEquals("schedule:7", userInput.getValue().getConversationId()); + } + + @Test + void shouldRejectSubMinuteCron() { + AgentSchedule input = schedule(); + input.setCronExpression("5 * * * * *"); + stubNoticeConfiguration(); + + assertThrows(IllegalArgumentException.class, () -> service.create(input)); + verify(scheduleDao, never()).saveAndFlush(any()); + } + + @Test + void shouldRequireExistingReceivers() { + AgentSchedule input = schedule(); + input.setReceiverIds(List.of(99L)); + + assertThrows(IllegalArgumentException.class, () -> service.create(input)); + } + + @Test + void shouldNotEnableMigratedScheduleWithoutReceiverConfiguration() { + AgentSchedule migrated = persistedSchedule(); + migrated.setEnabled(false); + migrated.setReceiverIds(List.of()); + when(scheduleDao.findById(7L)).thenReturn(Optional.of(migrated)); + + assertThrows(IllegalArgumentException.class, () -> service.toggle(7L, true)); + } + + @Test + void shouldSkipCronTriggerWhileRunIsActive() { + AgentSchedule schedule = persistedSchedule(); + when(scheduleDao.findById(7L)).thenReturn(Optional.of(schedule)); + when(runService.hasActiveRun(21L)).thenReturn(true); + + Optional run = service.claimCronRun(7L, System.currentTimeMillis()); + + assertFalse(run.isPresent()); + verify(runService, never()).createOrResumeRun(any(), any(), any()); + verify(scheduleDao).save(schedule); + } + + @Test + void shouldUsePlannedTimeInCronMessageId() { + AgentSchedule schedule = persistedSchedule(); + long plannedAt = schedule.getNextTriggerAt(); + AgentSession session = AgentSession.builder().id(21L).build(); + when(scheduleDao.findById(7L)).thenReturn(Optional.of(schedule)); + when(sessionService.findSession("21")).thenReturn(Optional.of(session)); + when(runService.createOrResumeRun(any(), any(), any())).thenAnswer(invocation -> { + UserInput input = invocation.getArgument(1); + return AgentRun.builder().messageId(input.getMessageId()).build(); + }); + + AgentRun run = service.claimCronRun(7L, System.currentTimeMillis()).orElseThrow(); + + assertEquals("schedule:7:cron:" + plannedAt, run.getMessageId()); + } + + private AgentSchedule schedule() { + return AgentSchedule.builder() + .name("Daily health inspection") + .instruction("Inspect all unhealthy monitors") + .cronExpression("0 0 9 * * *") + .enabled(true) + .receiverIds(List.of(10L)) + .build(); + } + + private AgentSchedule persistedSchedule() { + AgentSchedule schedule = schedule(); + schedule.setId(7L); + schedule.setSessionId(21L); + schedule.setNextTriggerAt(System.currentTimeMillis() - 1_000L); + return schedule; + } + + private void stubNoticeConfiguration() { + when(noticeConfigService.getReceiverById(10L)) + .thenReturn(NoticeReceiver.builder().id(10L).type((byte) 1).build()); + when(noticeConfigService.getDefaultNoticeTemplateByType((byte) 1)) + .thenReturn(NoticeTemplate.builder().type((byte) 1).build()); + } +} diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutorTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutorTest.java deleted file mode 100644 index cf1c9009ce..0000000000 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandExecutorTest.java +++ /dev/null @@ -1,73 +0,0 @@ -/* - * 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.gateway.schedule; - -import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - -import java.util.List; -import org.apache.hertzbeat.ai.gateway.application.GatewayCommand.InvokeCommand; -import org.apache.hertzbeat.ai.gateway.application.GatewayCommandRouter; -import org.apache.hertzbeat.ai.gateway.runtime.AgentRuntimeEntryType; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.ArgumentCaptor; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -/** - * Test case for {@link AgentScheduledCommandExecutor}. - */ -@ExtendWith(MockitoExtension.class) -class AgentScheduledCommandExecutorTest { - - @Mock - private AgentScheduledCommandService commandService; - - @Mock - private GatewayCommandRouter commandRouter; - - @Test - void shouldRouteDueCommandThroughGatewayRuntime() { - AgentScheduledCommand command = AgentScheduledCommand.builder() - .id(7L) - .channel("webui") - .conversationId("scheduled-conversation") - .actorType("user") - .actorId("admin") - .actorRoles("[\"admin\"]") - .message("Inspect all unhealthy monitors") - .enabled(true) - .build(); - when(commandService.findDueCommands()).thenReturn(List.of(command)); - - new AgentScheduledCommandExecutor(commandService, commandRouter).executeDueCommands(); - - ArgumentCaptor captor = ArgumentCaptor.forClass(InvokeCommand.class); - verify(commandRouter).handle(captor.capture()); - InvokeCommand invoke = captor.getValue(); - assertEquals(AgentRuntimeEntryType.SCHEDULE_TRIGGER, invoke.entryType()); - assertEquals("scheduled-conversation", invoke.userInput().getConversationId()); - assertEquals("Inspect all unhealthy monitors", invoke.userInput().getMessage().getText()); - assertEquals("admin", invoke.envelope().getActor().getId()); - assertEquals(org.apache.hertzbeat.ai.gateway.contract.AgentResponseLanguage.systemDefault(), - invoke.envelope().getPreferredLanguage()); - verify(commandService).completeExecution(command); - } -} diff --git a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandServiceTest.java b/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandServiceTest.java deleted file mode 100644 index e60cff65ad..0000000000 --- a/hertzbeat-ai-gateway/src/test/java/org/apache/hertzbeat/ai/gateway/schedule/AgentScheduledCommandServiceTest.java +++ /dev/null @@ -1,80 +0,0 @@ -/* - * 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.gateway.schedule; - -import static org.junit.jupiter.api.Assertions.assertNotNull; -import static org.junit.jupiter.api.Assertions.assertNull; -import static org.junit.jupiter.api.Assertions.assertThrows; -import static org.mockito.Mockito.when; - -import java.util.Optional; -import org.junit.jupiter.api.Test; -import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.Mock; -import org.mockito.junit.jupiter.MockitoExtension; - -/** - * Test case for {@link AgentScheduledCommandService}. - */ -@ExtendWith(MockitoExtension.class) -class AgentScheduledCommandServiceTest { - - @Mock - private AgentScheduledCommandDao commandDao; - - @Test - void shouldCalculateNextExecutionWhenCreatingCommand() { - AgentScheduledCommand command = AgentScheduledCommand.builder() - .sessionId(1L) - .message("Inspect monitor health") - .cronExpression("0 0 9 * * ?") - .build(); - when(commandDao.save(command)).thenReturn(command); - - AgentScheduledCommand result = service().create(command); - - assertNotNull(result.getNextRunTime()); - } - - @Test - void shouldNotScheduleDisabledCommand() { - AgentScheduledCommand command = AgentScheduledCommand.builder() - .sessionId(1L) - .message("Inspect monitor health") - .cronExpression("0 0 9 * * ?") - .enabled(false) - .build(); - when(commandDao.save(command)).thenReturn(command); - - AgentScheduledCommand result = service().create(command); - - assertNull(result.getNextRunTime()); - } - - @Test - void shouldRejectCommandOwnedByDifferentSession() { - AgentScheduledCommand command = AgentScheduledCommand.builder().id(2L).sessionId(1L).build(); - when(commandDao.findById(2L)).thenReturn(Optional.of(command)); - - assertThrows(IllegalArgumentException.class, () -> service().getOwned(2L, 3L)); - } - - private AgentScheduledCommandService service() { - return new AgentScheduledCommandService(commandDao); - } -} diff --git a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentRun.java b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentRun.java index 770c4bde8f..b6660f3b30 100644 --- a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentRun.java +++ b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentRun.java @@ -68,6 +68,9 @@ public class AgentRun { @Column(name = "message_id", nullable = false, length = 128) private String messageId; + @Column(name = "entry_type", nullable = false, length = 32) + private String entryType; + @Column(name = "target_monitor_id") private Long targetMonitorId; diff --git a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentSession.java b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentSession.java index 46ee566ec8..e07df8c0a4 100644 --- a/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentSession.java +++ b/hertzbeat-common-spring/src/main/java/org/apache/hertzbeat/common/entity/agent/AgentSession.java @@ -69,6 +69,9 @@ public class AgentSession { @Column(length = 64) private String channel; + @Column(name = "origin_entry_type", nullable = false, length = 32) + private String originEntryType; + @Column(name = "conversation_id", length = 256) private String conversationId; diff --git a/hertzbeat-e2e/hertzbeat-observability-e2e/src/test/resources/sureness.yml b/hertzbeat-e2e/hertzbeat-observability-e2e/src/test/resources/sureness.yml index e712ffe4b7..031685e19b 100644 --- a/hertzbeat-e2e/hertzbeat-observability-e2e/src/test/resources/sureness.yml +++ b/hertzbeat-e2e/hertzbeat-observability-e2e/src/test/resources/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/hertzbeat-manager/src/test/resources/sureness.yml b/hertzbeat-manager/src/test/resources/sureness.yml index 7707e4955a..abf96000f9 100644 --- a/hertzbeat-manager/src/test/resources/sureness.yml +++ b/hertzbeat-manager/src/test/resources/sureness.yml @@ -158,6 +158,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/hertzbeat-startup/src/main/java/org/apache/hertzbeat/startup/runtime/StartupRuntimeComponentBoundary.java b/hertzbeat-startup/src/main/java/org/apache/hertzbeat/startup/runtime/StartupRuntimeComponentBoundary.java index a117b39600..99ae140c72 100644 --- a/hertzbeat-startup/src/main/java/org/apache/hertzbeat/startup/runtime/StartupRuntimeComponentBoundary.java +++ b/hertzbeat-startup/src/main/java/org/apache/hertzbeat/startup/runtime/StartupRuntimeComponentBoundary.java @@ -33,7 +33,7 @@ final class StartupRuntimeComponentBoundary implements BeanDefinitionRegistryPos private static final String AGENT_GATEWAY_RESOURCE = "org/apache/hertzbeat/ai/gateway"; private static final Set AGENT_GATEWAY_INFRASTRUCTURE = Set.of( "agentGatewayRuntimeConfiguration", "agentRunDao", "agentSessionDao", - "agentTranscriptEntryDao", "agentScheduledCommandDao", "agentToolCallDao"); + "agentTranscriptEntryDao", "agentScheduleDao", "agentToolCallDao"); @Override public int getOrder() { diff --git a/hertzbeat-startup/src/main/resources/db/migration/h2/V200__create_entity_foundation.sql b/hertzbeat-startup/src/main/resources/db/migration/h2/V200__create_entity_foundation.sql index 941db497b3..55169f255f 100644 --- a/hertzbeat-startup/src/main/resources/db/migration/h2/V200__create_entity_foundation.sql +++ b/hertzbeat-startup/src/main/resources/db/migration/h2/V200__create_entity_foundation.sql @@ -304,6 +304,7 @@ CREATE TABLE IF NOT EXISTS hzb_agent_session ( session_uid VARCHAR(64) NOT NULL, session_key VARCHAR(128) NOT NULL, channel VARCHAR(64), + origin_entry_type VARCHAR(32) NOT NULL, conversation_id VARCHAR(256), actor_type VARCHAR(64), actor_id VARCHAR(128), @@ -323,6 +324,7 @@ CREATE TABLE IF NOT EXISTS hzb_agent_run ( run_uid VARCHAR(64) NOT NULL, session_id BIGINT NOT NULL, message_id VARCHAR(128) NOT NULL, + entry_type VARCHAR(32) NOT NULL, target_monitor_id BIGINT, target_alert_id BIGINT, target_collector VARCHAR(128), @@ -393,26 +395,27 @@ CREATE INDEX IF NOT EXISTS idx_agent_transcript_run CREATE INDEX IF NOT EXISTS idx_agent_transcript_checkpoint ON hzb_agent_transcript_entry(session_id, message_role, session_sequence); -CREATE TABLE IF NOT EXISTS hzb_agent_scheduled_command ( +CREATE TABLE IF NOT EXISTS hzb_agent_schedule ( id BIGINT AUTO_INCREMENT PRIMARY KEY, - session_id BIGINT NOT NULL, - channel VARCHAR(64) NOT NULL, - conversation_id VARCHAR(256) NOT NULL, - actor_type VARCHAR(64) NOT NULL, - actor_id VARCHAR(128) NOT NULL, - actor_roles VARCHAR(1024) NOT NULL, - message VARCHAR(4096) NOT NULL, + name VARCHAR(128) NOT NULL, + instruction VARCHAR(4096) NOT NULL, cron_expression VARCHAR(64) NOT NULL, - enabled BOOLEAN NOT NULL DEFAULT TRUE, - last_run_time TIMESTAMP, - next_run_time TIMESTAMP, + enabled BOOLEAN DEFAULT TRUE NOT NULL, + session_id BIGINT, + receiver_ids VARCHAR(2048) NOT NULL, + template_id BIGINT, + created_from_session_uid VARCHAR(64), + last_trigger_at BIGINT, + next_trigger_at BIGINT, + creator VARCHAR(64), + modifier VARCHAR(64), gmt_create TIMESTAMP DEFAULT CURRENT_TIMESTAMP, gmt_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); -CREATE INDEX IF NOT EXISTS idx_agent_scheduled_command_session - ON hzb_agent_scheduled_command(session_id); -CREATE INDEX IF NOT EXISTS idx_agent_scheduled_command_due - ON hzb_agent_scheduled_command(enabled, next_run_time); +CREATE INDEX IF NOT EXISTS idx_agent_schedule_due + ON hzb_agent_schedule(enabled, next_trigger_at); +CREATE INDEX IF NOT EXISTS idx_agent_schedule_session + ON hzb_agent_schedule(session_id); CREATE TABLE IF NOT EXISTS hzb_alert_analysis_policy ( id BIGINT AUTO_INCREMENT PRIMARY KEY, diff --git a/hertzbeat-startup/src/main/resources/db/migration/mysql/V200__create_entity_foundation.sql b/hertzbeat-startup/src/main/resources/db/migration/mysql/V200__create_entity_foundation.sql index 4209b37d68..29bc48189b 100644 --- a/hertzbeat-startup/src/main/resources/db/migration/mysql/V200__create_entity_foundation.sql +++ b/hertzbeat-startup/src/main/resources/db/migration/mysql/V200__create_entity_foundation.sql @@ -264,7 +264,8 @@ CREATE TABLE hzb_agent_session ( id BIGINT AUTO_INCREMENT PRIMARY KEY, session_uid VARCHAR(64) NOT NULL, session_key VARCHAR(128) NOT NULL, - channel VARCHAR(64), conversation_id VARCHAR(256), actor_type VARCHAR(64), + channel VARCHAR(64), origin_entry_type VARCHAR(32) NOT NULL, + conversation_id VARCHAR(256), actor_type VARCHAR(64), actor_id VARCHAR(128), actor_roles VARCHAR(1024), status VARCHAR(32), title VARCHAR(256), transcript_sequence BIGINT NOT NULL DEFAULT 0, gmt_create DATETIME DEFAULT CURRENT_TIMESTAMP, @@ -279,6 +280,7 @@ CREATE TABLE hzb_agent_run ( run_uid VARCHAR(64) NOT NULL, session_id BIGINT NOT NULL, message_id VARCHAR(128) NOT NULL, + entry_type VARCHAR(32) NOT NULL, target_monitor_id BIGINT, target_alert_id BIGINT, target_collector VARCHAR(128), target_context_json TEXT, status VARCHAR(32) NOT NULL, result_summary TEXT, error_message VARCHAR(1024), @@ -321,17 +323,17 @@ CREATE TABLE hzb_agent_transcript_entry ( KEY idx_agent_transcript_checkpoint (session_id, message_role, session_sequence) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci; -CREATE TABLE hzb_agent_scheduled_command ( +CREATE TABLE hzb_agent_schedule ( id BIGINT AUTO_INCREMENT PRIMARY KEY, - session_id BIGINT NOT NULL, channel VARCHAR(64) NOT NULL, - conversation_id VARCHAR(256) NOT NULL, actor_type VARCHAR(64) NOT NULL, - actor_id VARCHAR(128) NOT NULL, actor_roles VARCHAR(1024) NOT NULL, - message VARCHAR(4096) NOT NULL, cron_expression VARCHAR(64) NOT NULL, - enabled BOOLEAN NOT NULL DEFAULT TRUE, last_run_time DATETIME, next_run_time DATETIME, + name VARCHAR(128) NOT NULL, instruction VARCHAR(4096) NOT NULL, + cron_expression VARCHAR(64) NOT NULL, enabled BOOLEAN NOT NULL DEFAULT TRUE, + session_id BIGINT, receiver_ids VARCHAR(2048) NOT NULL, template_id BIGINT, + created_from_session_uid VARCHAR(64), last_trigger_at BIGINT, next_trigger_at BIGINT, + creator VARCHAR(64), modifier VARCHAR(64), gmt_create DATETIME DEFAULT CURRENT_TIMESTAMP, gmt_update DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, - KEY idx_agent_scheduled_command_session (session_id), - KEY idx_agent_scheduled_command_due (enabled, next_run_time) + KEY idx_agent_schedule_due (enabled, next_trigger_at), + KEY idx_agent_schedule_session (session_id) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_general_ci; CREATE TABLE hzb_alert_analysis_policy ( diff --git a/hertzbeat-startup/src/main/resources/db/migration/postgresql/V200__create_entity_foundation.sql b/hertzbeat-startup/src/main/resources/db/migration/postgresql/V200__create_entity_foundation.sql index f6dc5a735b..8c671bd1df 100644 --- a/hertzbeat-startup/src/main/resources/db/migration/postgresql/V200__create_entity_foundation.sql +++ b/hertzbeat-startup/src/main/resources/db/migration/postgresql/V200__create_entity_foundation.sql @@ -302,7 +302,8 @@ CREATE TABLE hzb_agent_session ( id BIGSERIAL PRIMARY KEY, session_uid VARCHAR(64) NOT NULL, session_key VARCHAR(128) NOT NULL, - channel VARCHAR(64), conversation_id VARCHAR(256), actor_type VARCHAR(64), + channel VARCHAR(64), origin_entry_type VARCHAR(32) NOT NULL, + conversation_id VARCHAR(256), actor_type VARCHAR(64), actor_id VARCHAR(128), actor_roles VARCHAR(1024), status VARCHAR(32), title VARCHAR(256), transcript_sequence BIGINT NOT NULL DEFAULT 0, gmt_create TIMESTAMP DEFAULT CURRENT_TIMESTAMP, @@ -315,6 +316,7 @@ CREATE INDEX idx_agent_session_owner ON hzb_agent_session(channel, actor_type, a CREATE TABLE hzb_agent_run ( id BIGSERIAL PRIMARY KEY, run_uid VARCHAR(64) NOT NULL, session_id BIGINT NOT NULL, message_id VARCHAR(128) NOT NULL, + entry_type VARCHAR(32) NOT NULL, target_monitor_id BIGINT, target_alert_id BIGINT, target_collector VARCHAR(128), target_context_json TEXT, status VARCHAR(32) NOT NULL, result_summary TEXT, error_message VARCHAR(1024), @@ -358,17 +360,17 @@ CREATE INDEX idx_agent_transcript_run ON hzb_agent_transcript_entry(run_id, sess CREATE INDEX idx_agent_transcript_checkpoint ON hzb_agent_transcript_entry(session_id, message_role, session_sequence); -CREATE TABLE hzb_agent_scheduled_command ( +CREATE TABLE hzb_agent_schedule ( id BIGSERIAL PRIMARY KEY, - session_id BIGINT NOT NULL, channel VARCHAR(64) NOT NULL, - conversation_id VARCHAR(256) NOT NULL, actor_type VARCHAR(64) NOT NULL, - actor_id VARCHAR(128) NOT NULL, actor_roles VARCHAR(1024) NOT NULL, - message VARCHAR(4096) NOT NULL, cron_expression VARCHAR(64) NOT NULL, - enabled BOOLEAN NOT NULL DEFAULT TRUE, last_run_time TIMESTAMP, next_run_time TIMESTAMP, + name VARCHAR(128) NOT NULL, instruction VARCHAR(4096) NOT NULL, + cron_expression VARCHAR(64) NOT NULL, enabled BOOLEAN NOT NULL DEFAULT TRUE, + session_id BIGINT, receiver_ids VARCHAR(2048) NOT NULL, template_id BIGINT, + created_from_session_uid VARCHAR(64), last_trigger_at BIGINT, next_trigger_at BIGINT, + creator VARCHAR(64), modifier VARCHAR(64), gmt_create TIMESTAMP DEFAULT CURRENT_TIMESTAMP, gmt_update TIMESTAMP DEFAULT CURRENT_TIMESTAMP ); -CREATE INDEX idx_agent_scheduled_command_session ON hzb_agent_scheduled_command(session_id); -CREATE INDEX idx_agent_scheduled_command_due ON hzb_agent_scheduled_command(enabled, next_run_time); +CREATE INDEX idx_agent_schedule_due ON hzb_agent_schedule(enabled, next_trigger_at); +CREATE INDEX idx_agent_schedule_session ON hzb_agent_schedule(session_id); CREATE TABLE hzb_alert_analysis_policy ( id BIGSERIAL PRIMARY KEY, diff --git a/hertzbeat-startup/src/main/resources/sureness.yml b/hertzbeat-startup/src/main/resources/sureness.yml index db809d7032..b7736ce973 100644 --- a/hertzbeat-startup/src/main/resources/sureness.yml +++ b/hertzbeat-startup/src/main/resources/sureness.yml @@ -167,6 +167,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/AgentGatewayAuthorizationConfigTest.java b/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/AgentGatewayAuthorizationConfigTest.java index 3ff073c05a..64f4a817de 100644 --- a/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/AgentGatewayAuthorizationConfigTest.java +++ b/hertzbeat-startup/src/test/java/org/apache/hertzbeat/startup/AgentGatewayAuthorizationConfigTest.java @@ -46,6 +46,11 @@ class AgentGatewayAuthorizationConfigTest { " - /api/agent/model-providers/**===post===[admin]", " - /api/agent/model-providers/**===put===[admin]", " - /api/agent/model-providers/**===delete===[admin]", + " - /api/agent/schedules/**===get===[admin]", + " - /api/agent/schedules/**===post===[admin]", + " - /api/agent/schedules/**===put===[admin]", + " - /api/agent/schedules/**===patch===[admin]", + " - /api/agent/schedules/**===delete===[admin]", " - /api/agent/alert-analysis/**===get===[admin]", " - /api/agent/sessions===get===[admin,user]", " - /api/agent/sessions/**===get===[admin,user]", diff --git a/script/docker-compose/hertzbeat-mysql-iotdb/conf/sureness.yml b/script/docker-compose/hertzbeat-mysql-iotdb/conf/sureness.yml index 2be16ad17e..fd253901f5 100644 --- a/script/docker-compose/hertzbeat-mysql-iotdb/conf/sureness.yml +++ b/script/docker-compose/hertzbeat-mysql-iotdb/conf/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/script/docker-compose/hertzbeat-mysql-tdengine/conf/sureness.yml b/script/docker-compose/hertzbeat-mysql-tdengine/conf/sureness.yml index 2be16ad17e..fd253901f5 100644 --- a/script/docker-compose/hertzbeat-mysql-tdengine/conf/sureness.yml +++ b/script/docker-compose/hertzbeat-mysql-tdengine/conf/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/script/docker-compose/hertzbeat-mysql-victoria-metrics/conf/sureness.yml b/script/docker-compose/hertzbeat-mysql-victoria-metrics/conf/sureness.yml index 2be16ad17e..fd253901f5 100644 --- a/script/docker-compose/hertzbeat-mysql-victoria-metrics/conf/sureness.yml +++ b/script/docker-compose/hertzbeat-mysql-victoria-metrics/conf/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/script/docker-compose/hertzbeat-postgresql-greptimedb/conf/sureness.yml b/script/docker-compose/hertzbeat-postgresql-greptimedb/conf/sureness.yml index 2be16ad17e..fd253901f5 100644 --- a/script/docker-compose/hertzbeat-postgresql-greptimedb/conf/sureness.yml +++ b/script/docker-compose/hertzbeat-postgresql-greptimedb/conf/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/script/docker-compose/hertzbeat-postgresql-victoria-metrics/conf/sureness.yml b/script/docker-compose/hertzbeat-postgresql-victoria-metrics/conf/sureness.yml index 2be16ad17e..fd253901f5 100644 --- a/script/docker-compose/hertzbeat-postgresql-victoria-metrics/conf/sureness.yml +++ b/script/docker-compose/hertzbeat-postgresql-victoria-metrics/conf/sureness.yml @@ -152,6 +152,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user] diff --git a/script/sureness.yml b/script/sureness.yml index c0e11ca0ad..7f795921c1 100644 --- a/script/sureness.yml +++ b/script/sureness.yml @@ -157,6 +157,11 @@ resourceRole: - /api/agent/model-providers/**===post===[admin] - /api/agent/model-providers/**===put===[admin] - /api/agent/model-providers/**===delete===[admin] + - /api/agent/schedules/**===get===[admin] + - /api/agent/schedules/**===post===[admin] + - /api/agent/schedules/**===put===[admin] + - /api/agent/schedules/**===patch===[admin] + - /api/agent/schedules/**===delete===[admin] - /api/agent/alert-analysis/**===get===[admin] - /api/agent/sessions===get===[admin,user] - /api/agent/sessions/**===get===[admin,user]