feat: 支持 Agent 运行刷新恢复
- 解耦 Runtime 与浏览器 SSE 订阅并增加 Redis 游标日志 - 支持正式聊天与草稿试运行刷新重连和权威终态恢复 - 完善显式取消、owner 丢失、容量限制与故障测试
This commit is contained in:
@@ -72,6 +72,21 @@ public class AgentRuntimeProperties {
|
||||
*/
|
||||
private Duration asyncToolTaskTtl = Duration.ofHours(24);
|
||||
|
||||
/**
|
||||
* AG-UI 可重连运行事件保留时间。
|
||||
*/
|
||||
private Duration aguiRunRetention = Duration.ofHours(24);
|
||||
|
||||
/**
|
||||
* AG-UI 事件日志批量刷入 Redis 的间隔。
|
||||
*/
|
||||
private Duration aguiJournalFlushInterval = Duration.ofMillis(50);
|
||||
|
||||
/**
|
||||
* AG-UI 重连订阅轮询 Redis 的间隔。
|
||||
*/
|
||||
private Duration aguiReplayPollInterval = Duration.ofMillis(100);
|
||||
|
||||
/**
|
||||
* 获取 Redis 热态 session 缓存 TTL。
|
||||
*
|
||||
@@ -90,6 +105,60 @@ public class AgentRuntimeProperties {
|
||||
this.sessionCacheTtl = sessionCacheTtl == null ? Duration.ofHours(24) : sessionCacheTtl;
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取 AG-UI 可重连运行事件保留时间。
|
||||
*
|
||||
* @return 事件保留时间
|
||||
*/
|
||||
public Duration getAguiRunRetention() {
|
||||
return aguiRunRetention;
|
||||
}
|
||||
|
||||
/**
|
||||
* 设置 AG-UI 可重连运行事件保留时间。
|
||||
*
|
||||
* @param aguiRunRetention 事件保留时间
|
||||
*/
|
||||
public void setAguiRunRetention(Duration aguiRunRetention) {
|
||||
this.aguiRunRetention = positiveDuration(aguiRunRetention, Duration.ofHours(24));
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取 AG-UI 事件日志批量刷入 Redis 的间隔。
|
||||
*
|
||||
* @return 刷盘间隔
|
||||
*/
|
||||
public Duration getAguiJournalFlushInterval() {
|
||||
return aguiJournalFlushInterval;
|
||||
}
|
||||
|
||||
/**
|
||||
* 设置 AG-UI 事件日志批量刷入 Redis 的间隔。
|
||||
*
|
||||
* @param aguiJournalFlushInterval 刷盘间隔
|
||||
*/
|
||||
public void setAguiJournalFlushInterval(Duration aguiJournalFlushInterval) {
|
||||
this.aguiJournalFlushInterval = positiveDuration(aguiJournalFlushInterval, Duration.ofMillis(50));
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取 AG-UI 重连订阅轮询 Redis 的间隔。
|
||||
*
|
||||
* @return 轮询间隔
|
||||
*/
|
||||
public Duration getAguiReplayPollInterval() {
|
||||
return aguiReplayPollInterval;
|
||||
}
|
||||
|
||||
/**
|
||||
* 设置 AG-UI 重连订阅轮询 Redis 的间隔。
|
||||
*
|
||||
* @param aguiReplayPollInterval 轮询间隔
|
||||
*/
|
||||
public void setAguiReplayPollInterval(Duration aguiReplayPollInterval) {
|
||||
this.aguiReplayPollInterval = positiveDuration(aguiReplayPollInterval, Duration.ofMillis(100));
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取当前 Agent 运行实例 ID。
|
||||
*
|
||||
@@ -281,6 +350,10 @@ public class AgentRuntimeProperties {
|
||||
this.asyncToolTaskTtl = asyncToolTaskTtl == null ? Duration.ofHours(24) : asyncToolTaskTtl;
|
||||
}
|
||||
|
||||
private static Duration positiveDuration(Duration value, Duration fallback) {
|
||||
return value == null || value.isZero() || value.isNegative() ? fallback : value;
|
||||
}
|
||||
|
||||
private static String defaultInstanceId() {
|
||||
String envInstanceId = System.getenv("EASYFLOW_INSTANCE_ID");
|
||||
if (StringUtils.hasText(envInstanceId)) {
|
||||
|
||||
@@ -23,5 +23,10 @@ public enum AgentRuntimeCommandAction {
|
||||
/**
|
||||
* 取消指定 Agent 在目标节点上的全部运行。
|
||||
*/
|
||||
CANCEL_AGENT
|
||||
CANCEL_AGENT,
|
||||
|
||||
/**
|
||||
* 取消指定的单次运行。
|
||||
*/
|
||||
CANCEL_RUN
|
||||
}
|
||||
|
||||
@@ -108,6 +108,9 @@ public class AgentRuntimeCommandConsumer implements MQConsumerHandler {
|
||||
command.getRequestId(), command.getResumeToken(), command.getReason());
|
||||
} else if (command.getAction() == AgentRuntimeCommandAction.CANCEL_AGENT) {
|
||||
agentRunService.cancelAgentLocal(command.getAgentId());
|
||||
} else if (command.getAction() == AgentRuntimeCommandAction.CANCEL_RUN) {
|
||||
agentRunService.cancelRunLocal(
|
||||
command.getRequestId(), command.getUserId(), command.getReason());
|
||||
} else {
|
||||
markFailureQuietly(command, new IllegalArgumentException("不支持的 Agent 远程运行命令"));
|
||||
LOG.warn("跳过不支持的 Agent 远程运行命令: messageId={}, commandId={}, action={}",
|
||||
|
||||
@@ -175,6 +175,21 @@ public class AgentRuntimeCommandProducer {
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 投递远程单次运行取消命令。
|
||||
*
|
||||
* @param targetNodeId 目标节点 ID
|
||||
* @param requestId 内部请求 ID
|
||||
* @param userId 当前用户 ID
|
||||
* @param reason 取消原因
|
||||
*/
|
||||
public void sendCancelRun(String targetNodeId, String requestId, String userId, String reason) {
|
||||
sendAndWait(
|
||||
targetNodeId, requestId, null, null, null,
|
||||
AgentRuntimeCommandAction.CANCEL_RUN, reason, null, userId
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* 投递远程运行命令并等待目标节点确认。
|
||||
*
|
||||
|
||||
@@ -93,9 +93,13 @@ public class AgentRunRegistry {
|
||||
*/
|
||||
public void bindSubscription(String requestId, Disposable subscription) {
|
||||
AgentRunContext context = runs.get(requestId);
|
||||
if (context != null) {
|
||||
context.setSubscription(subscription);
|
||||
if (context == null) {
|
||||
if (subscription != null && !subscription.isDisposed()) {
|
||||
subscription.dispose();
|
||||
}
|
||||
return;
|
||||
}
|
||||
context.setSubscription(subscription);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -152,6 +156,31 @@ public class AgentRunRegistry {
|
||||
remove(requestId);
|
||||
}
|
||||
|
||||
/**
|
||||
* 显式取消当前节点上的单次运行，并通过原事件处理链生成规范取消终态。
|
||||
*
|
||||
* @param requestId 内部请求 ID
|
||||
* @param userId 当前用户 ID
|
||||
* @param reason 取消原因
|
||||
*/
|
||||
public void cancelRun(String requestId, String userId, String reason) {
|
||||
if (requestId == null || requestId.isBlank()) {
|
||||
throw new BusinessException("Agent 运行请求 ID 不能为空");
|
||||
}
|
||||
AgentRunContext context = runs.get(requestId);
|
||||
if (context == null) {
|
||||
// 取消命令允许因自然终态、重复投递或跨节点竞态而幂等到达。
|
||||
return;
|
||||
}
|
||||
assertOwner(requestId, userId);
|
||||
context.cancel();
|
||||
AgentRuntimeEvent event = AgentRuntimeEvent.of(
|
||||
com.easyagents.agent.runtime.event.AgentRuntimeEventType.CANCELLED);
|
||||
event.getPayload().put("reason",
|
||||
reason == null || reason.isBlank() ? "用户已停止生成" : reason);
|
||||
context.eventConsumer().accept(event);
|
||||
}
|
||||
|
||||
/**
|
||||
* 取消当前节点上指定 Agent 的全部活跃运行。
|
||||
*
|
||||
@@ -314,34 +343,80 @@ public class AgentRunRegistry {
|
||||
if (requestId == null) {
|
||||
return;
|
||||
}
|
||||
RuntimeException cleanupFailure = null;
|
||||
AgentRunContext context = runs.remove(requestId);
|
||||
if (context != null) {
|
||||
sessionRuns.remove(context.sessionId(), requestId);
|
||||
context.releaseLock();
|
||||
try {
|
||||
context.releaseLock();
|
||||
} catch (RuntimeException exception) {
|
||||
cleanupFailure = mergeCleanupFailure(
|
||||
cleanupFailure, exception, requestId, "release distributed lock");
|
||||
}
|
||||
// Redis 解锁失败也不能阻止底层模型、工具与订阅资源释放。
|
||||
context.closeRuntime();
|
||||
}
|
||||
owners.remove(requestId);
|
||||
Set<String> tokens = requestTokens.remove(requestId);
|
||||
if (tokens != null) {
|
||||
tokens.forEach(token -> {
|
||||
for (String token : tokens) {
|
||||
resumeTokenIndex.remove(token);
|
||||
if (routeRegistry != null) {
|
||||
routeRegistry.removeResumeToken(token);
|
||||
try {
|
||||
routeRegistry.removeResumeToken(token);
|
||||
} catch (RuntimeException exception) {
|
||||
cleanupFailure = mergeCleanupFailure(
|
||||
cleanupFailure, exception, requestId, "remove resume token route");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
Set<String> approvals = requestApprovals.remove(requestId);
|
||||
if (approvals != null) {
|
||||
approvals.forEach(approvalId -> {
|
||||
for (String approvalId : approvals) {
|
||||
approvalTargets.remove(approvalId);
|
||||
if (routeRegistry != null) {
|
||||
routeRegistry.removeApproval(approvalId);
|
||||
try {
|
||||
routeRegistry.removeApproval(approvalId);
|
||||
} catch (RuntimeException exception) {
|
||||
cleanupFailure = mergeCleanupFailure(
|
||||
cleanupFailure, exception, requestId, "remove approval route");
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
if (routeRegistry != null) {
|
||||
routeRegistry.removeRun(requestId);
|
||||
try {
|
||||
routeRegistry.removeRun(requestId);
|
||||
} catch (RuntimeException exception) {
|
||||
cleanupFailure = mergeCleanupFailure(
|
||||
cleanupFailure, exception, requestId, "remove run route");
|
||||
}
|
||||
}
|
||||
if (cleanupFailure != null) {
|
||||
throw cleanupFailure;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 合并运行清理异常，同时保留第一个异常作为调用方可见的根因。
|
||||
*
|
||||
* @param current 当前已记录异常
|
||||
* @param next 本次清理异常
|
||||
* @param requestId 请求 ID
|
||||
* @param phase 失败阶段
|
||||
* @return 合并后的首异常
|
||||
*/
|
||||
private RuntimeException mergeCleanupFailure(RuntimeException current,
|
||||
RuntimeException next,
|
||||
String requestId,
|
||||
String phase) {
|
||||
LOG.warn("Agent runtime cleanup failed, requestId={}, phase={}", requestId, phase, next);
|
||||
if (current == null) {
|
||||
return next;
|
||||
}
|
||||
current.addSuppressed(next);
|
||||
return current;
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -684,10 +759,18 @@ public class AgentRunRegistry {
|
||||
if (subscription == null) {
|
||||
return;
|
||||
}
|
||||
if (finished.get()) {
|
||||
subscription.dispose();
|
||||
return;
|
||||
}
|
||||
Disposable previous = this.subscription.getAndSet(subscription);
|
||||
if (previous != null && !previous.isDisposed()) {
|
||||
previous.dispose();
|
||||
}
|
||||
if (finished.get() && this.subscription.compareAndSet(subscription, null)
|
||||
&& !subscription.isDisposed()) {
|
||||
subscription.dispose();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
|
||||
@@ -32,8 +32,15 @@ import tech.easyflow.agent.distributed.AgentRuntimeRoute;
|
||||
import tech.easyflow.agent.distributed.AgentRuntimeRouteRegistry;
|
||||
import tech.easyflow.agent.runtime.event.AgentRunEventRecorder;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunInputMapper;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunDescriptor;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunJournal;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunStatus;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunStatusView;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunStore;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunSubscriptionService;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiWireContext;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiHitlResolveRequest;
|
||||
import tech.easyflow.agent.runtime.agui.ResumableAguiSseEmitter;
|
||||
import tech.easyflow.agent.runtime.artifact.AgentArtifactService;
|
||||
import tech.easyflow.agent.runtime.hitl.AgentHitlPendingService;
|
||||
import tech.easyflow.agent.runtime.hitl.ToolApprovalInputProjection;
|
||||
@@ -154,6 +161,12 @@ public class AgentRunService {
|
||||
@Resource
|
||||
private AgentAguiRunInputMapper agentAguiRunInputMapper;
|
||||
@Resource
|
||||
private AgentAguiRunStore agentAguiRunStore;
|
||||
@Resource
|
||||
private AgentAguiRunJournal agentAguiRunJournal;
|
||||
@Resource
|
||||
private AgentAguiRunSubscriptionService agentAguiRunSubscriptionService;
|
||||
@Resource
|
||||
private AgentBuiltinToolsConfigResolver agentBuiltinToolsConfigResolver;
|
||||
@Resource
|
||||
private AgentArtifactService agentArtifactService;
|
||||
@@ -217,10 +230,11 @@ public class AgentRunService {
|
||||
}
|
||||
applyFormalSessionTitle(chatContext, titlePrompt, existingSession);
|
||||
// 执行对话
|
||||
return run(agent, chatRequest.getPrompt(), mediaUploads, documentUploads,
|
||||
AgentRunOutput runOutput = createRunOutput(
|
||||
wireContext, requestId, account, agent, sessionId.toString(), false);
|
||||
return runAguiAware(agent, chatRequest.getPrompt(), mediaUploads, documentUploads,
|
||||
account, requestId, traceId, sessionId.toString(),
|
||||
ASSISTANT_CODE, chatContext, true, easyFlowAgentSessionStore,
|
||||
createRunOutput(wireContext));
|
||||
ASSISTANT_CODE, chatContext, true, easyFlowAgentSessionStore, runOutput);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -336,10 +350,11 @@ public class AgentRunService {
|
||||
String titlePrompt = effectivePrompt(draftRequest.getPrompt(), !documentUploads.isEmpty(), !mediaUploads.isEmpty());
|
||||
ChatRuntimeContext chatContext = buildChatRuntimeContext(
|
||||
agent, chatSessionId, titlePrompt, account, DRAFT_ASSISTANT_CODE);
|
||||
return run(agent, draftRequest.getPrompt(), mediaUploads, documentUploads,
|
||||
AgentRunOutput runOutput = createRunOutput(
|
||||
wireContext, requestId, account, agent, runtimeSessionId, true);
|
||||
return runAguiAware(agent, draftRequest.getPrompt(), mediaUploads, documentUploads,
|
||||
account, requestId, traceId, runtimeSessionId,
|
||||
DRAFT_ASSISTANT_CODE, chatContext, false, draftAgentSessionStore,
|
||||
createRunOutput(wireContext));
|
||||
DRAFT_ASSISTANT_CODE, chatContext, false, draftAgentSessionStore, runOutput);
|
||||
}
|
||||
|
||||
private SseEmitter run(Agent agent,
|
||||
@@ -415,22 +430,110 @@ public class AgentRunService {
|
||||
submitted = true;
|
||||
return runOutput.emitter();
|
||||
} finally {
|
||||
// 释放锁
|
||||
if (!submitted && lockHandle != null) {
|
||||
lockHandle.release();
|
||||
if (!submitted) {
|
||||
try {
|
||||
if (lockHandle != null) {
|
||||
lockHandle.release();
|
||||
}
|
||||
} finally {
|
||||
if (!runOutput.cancelRunOnDisconnect()) {
|
||||
// 可重连运行已在编译前注册临时 owner；未提交到后台时必须同步移除。
|
||||
agentRuntimeRouteRegistry.removeRun(requestId);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private AgentRunOutput createRunOutput(AgentAguiWireContext wireContext) {
|
||||
private AgentRunOutput createRunOutput(AgentAguiWireContext wireContext,
|
||||
String requestId,
|
||||
LoginAccount account,
|
||||
Agent agent,
|
||||
String runtimeSessionId,
|
||||
boolean draft) {
|
||||
if (wireContext == null) {
|
||||
return new LegacyAgentRunOutput();
|
||||
}
|
||||
return new AguiAgentRunOutput(
|
||||
wireContext.threadId(),
|
||||
long now = System.currentTimeMillis();
|
||||
AgentAguiRunDescriptor descriptor = new AgentAguiRunDescriptor(
|
||||
wireContext.runId(),
|
||||
wireContext.userMessageId(),
|
||||
wireContext.userMessageContent());
|
||||
requestId,
|
||||
wireContext.threadId(),
|
||||
agent == null || agent.getId() == null ? "" : agent.getId().toString(),
|
||||
runtimeSessionId,
|
||||
account.getId().toString(),
|
||||
account.getTenantId().toString(),
|
||||
draft,
|
||||
AgentAguiRunStatus.RUNNING,
|
||||
0L,
|
||||
now,
|
||||
now);
|
||||
boolean storeCreated = false;
|
||||
try {
|
||||
// 在耗时编译和初始化前注册 owner，刷新恢复可据节点心跳区分“仍在启动”与 owner 丢失。
|
||||
agentRuntimeRouteRegistry.registerRun(requestId, descriptor.agentId());
|
||||
agentAguiRunStore.create(descriptor);
|
||||
storeCreated = true;
|
||||
SseEmitter subscriberEmitter = agentAguiRunSubscriptionService.subscribe(descriptor, 0L);
|
||||
return new AguiAgentRunOutput(
|
||||
wireContext.threadId(),
|
||||
wireContext.runId(),
|
||||
wireContext.userMessageId(),
|
||||
wireContext.userMessageContent(),
|
||||
new ResumableAguiSseEmitter(
|
||||
wireContext.threadId(), wireContext.runId(),
|
||||
agentAguiRunJournal, subscriberEmitter),
|
||||
false);
|
||||
} catch (RuntimeException exception) {
|
||||
if (storeCreated) {
|
||||
try {
|
||||
agentAguiRunJournal.fail(
|
||||
descriptor.threadId(), descriptor.runId(),
|
||||
"AGENT_RUN_START_FAILED", "Agent 运行启动失败");
|
||||
} catch (RuntimeException journalException) {
|
||||
exception.addSuppressed(journalException);
|
||||
}
|
||||
}
|
||||
try {
|
||||
agentRuntimeRouteRegistry.removeRun(requestId);
|
||||
} catch (RuntimeException routeException) {
|
||||
exception.addSuppressed(routeException);
|
||||
}
|
||||
throw exception;
|
||||
}
|
||||
}
|
||||
|
||||
private SseEmitter runAguiAware(Agent agent,
|
||||
String prompt,
|
||||
List<AgentMediaUploadRecord> mediaUploads,
|
||||
List<AgentDocumentAttachment> documentUploads,
|
||||
LoginAccount account,
|
||||
String requestId,
|
||||
String traceId,
|
||||
String runtimeSessionId,
|
||||
String assistantCode,
|
||||
ChatRuntimeContext chatContext,
|
||||
boolean persistChatlog,
|
||||
AgentSessionStore runtimeSessionStore,
|
||||
AgentRunOutput runOutput) {
|
||||
try {
|
||||
return run(agent, prompt, mediaUploads, documentUploads, account, requestId, traceId,
|
||||
runtimeSessionId, assistantCode, chatContext, persistChatlog, runtimeSessionStore, runOutput);
|
||||
} catch (RuntimeException exception) {
|
||||
if (!runOutput.cancelRunOnDisconnect()) {
|
||||
try {
|
||||
runOutput.completeWithError(exception);
|
||||
} catch (RuntimeException outputException) {
|
||||
exception.addSuppressed(outputException);
|
||||
}
|
||||
try {
|
||||
agentRuntimeRouteRegistry.removeRun(requestId);
|
||||
} catch (RuntimeException routeException) {
|
||||
exception.addSuppressed(routeException);
|
||||
}
|
||||
}
|
||||
throw exception;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -610,6 +713,88 @@ public class AgentRunService {
|
||||
account.getTenantId() == null ? null : account.getTenantId().toString());
|
||||
}
|
||||
|
||||
/**
|
||||
* 查询当前用户可访问的 AG-UI 运行状态。
|
||||
*
|
||||
* @param runId 客户端运行 ID
|
||||
* @return 运行公开状态
|
||||
*/
|
||||
public AgentAguiRunStatusView getAguiRunStatus(String runId) {
|
||||
LoginAccount account = requireCurrentLoginAccount();
|
||||
return AgentAguiRunStatusView.from(reconcileInterruptedRun(
|
||||
agentAguiRunStore.requireOwned(runId, account)));
|
||||
}
|
||||
|
||||
/**
|
||||
* 从指定游标订阅当前用户可访问的 AG-UI 运行事件。
|
||||
*
|
||||
* @param runId 客户端运行 ID
|
||||
* @param afterCursor 已消费的最后事件游标
|
||||
* @return 增量重放 SSE
|
||||
*/
|
||||
public SseEmitter subscribeAguiRun(String runId, long afterCursor) {
|
||||
LoginAccount account = requireCurrentLoginAccount();
|
||||
AgentAguiRunDescriptor descriptor = reconcileInterruptedRun(
|
||||
agentAguiRunStore.requireOwned(runId, account));
|
||||
return agentAguiRunSubscriptionService.subscribe(descriptor, afterCursor);
|
||||
}
|
||||
|
||||
/**
|
||||
* 显式取消当前用户拥有的单次 AG-UI 运行。
|
||||
*
|
||||
* @param runId 客户端运行 ID
|
||||
*/
|
||||
public void cancelAguiRun(String runId) {
|
||||
LoginAccount account = requireCurrentLoginAccount();
|
||||
AgentAguiRunDescriptor descriptor = reconcileInterruptedRun(
|
||||
agentAguiRunStore.requireOwned(runId, account));
|
||||
if (descriptor.status().isTerminal()) {
|
||||
return;
|
||||
}
|
||||
agentAguiRunStore.requestCancellation(runId);
|
||||
String userId = account.getId().toString();
|
||||
if (agentRunRegistry.get(descriptor.requestId()) != null) {
|
||||
cancelRunLocal(descriptor.requestId(), userId, "用户已停止生成");
|
||||
return;
|
||||
}
|
||||
AgentRuntimeRoute route = agentRuntimeRouteRegistry.findOwnerRoute(descriptor.requestId());
|
||||
if (route == null || route.getNodeId() == null || route.getNodeId().isBlank()) {
|
||||
// Runtime 可能仍在编译或初始化，启动路径会在注册前后消费 cancel intent。
|
||||
return;
|
||||
}
|
||||
agentRuntimeCommandProducer.sendCancelRun(
|
||||
route.getNodeId(), descriptor.requestId(), userId, "用户已停止生成");
|
||||
}
|
||||
|
||||
/**
|
||||
* 在当前节点显式取消指定运行。
|
||||
*
|
||||
* @param requestId 内部请求 ID
|
||||
* @param userId 当前用户 ID
|
||||
* @param reason 取消原因
|
||||
*/
|
||||
public void cancelRunLocal(String requestId, String userId, String reason) {
|
||||
agentRunRegistry.cancelRun(requestId, userId, reason);
|
||||
}
|
||||
|
||||
private AgentAguiRunDescriptor reconcileInterruptedRun(AgentAguiRunDescriptor descriptor) {
|
||||
if (descriptor == null || descriptor.status().isTerminal()) {
|
||||
return descriptor;
|
||||
}
|
||||
AgentRuntimeRoute route = agentRuntimeRouteRegistry.findOwnerRoute(descriptor.requestId());
|
||||
boolean stale = route == null;
|
||||
if (route != null && route.getBootId() != null && route.getNodeId() != null) {
|
||||
String currentBootId = agentRuntimeRouteRegistry.currentNodeBootId(route.getNodeId());
|
||||
stale = currentBootId == null || !route.getBootId().equals(currentBootId);
|
||||
}
|
||||
if (!stale || System.currentTimeMillis() - descriptor.updatedAt() < 90_000L) {
|
||||
return descriptor;
|
||||
}
|
||||
agentAguiRunStore.failOwnerLost(descriptor);
|
||||
AgentAguiRunDescriptor refreshed = agentAguiRunStore.find(descriptor.runId());
|
||||
return refreshed == null ? descriptor : refreshed;
|
||||
}
|
||||
|
||||
private void clearDraftSessionInternal(String sessionId, String userId, String tenantId) {
|
||||
agentRunRegistry.cancelSession(sessionId, userId);
|
||||
draftAgentSessionStore.delete(sessionId);
|
||||
@@ -942,10 +1127,20 @@ public class AgentRunService {
|
||||
} catch (Exception exception) {
|
||||
AgentRunRegistry.AgentRunContext runContext = agentRunRegistry.get(requestId);
|
||||
if (runContext != null) {
|
||||
runContext.cancel();
|
||||
agentRunRegistry.remove(requestId);
|
||||
try {
|
||||
runContext.cancel();
|
||||
} catch (RuntimeException cancelException) {
|
||||
LOG.warn("Cancel Agent runtime after startup failure failed, requestId={}",
|
||||
requestId, cancelException);
|
||||
}
|
||||
try {
|
||||
agentRunRegistry.remove(requestId);
|
||||
} catch (RuntimeException removeException) {
|
||||
LOG.warn("Remove Agent runtime after startup failure failed, requestId={}",
|
||||
requestId, removeException);
|
||||
}
|
||||
} else if (initialLockHandle != null) {
|
||||
initialLockHandle.release();
|
||||
releaseRunLockQuietly(initialLockHandle, requestId);
|
||||
}
|
||||
handleRuntimeError(
|
||||
exception,
|
||||
@@ -997,6 +1192,15 @@ public class AgentRunService {
|
||||
assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
AgentRunLock.Handle lockHandle = initialLockHandle;
|
||||
try {
|
||||
if (isAguiCancellationRequested(runOutput)) {
|
||||
handleRuntimeEvent(cancellationEvent("用户已停止生成"), requestId, runOutput,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser,
|
||||
chatContext, finished, persistChatlog);
|
||||
if (lockHandle != null) {
|
||||
releaseRunLockQuietly(lockHandle, requestId);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (persistChatlog) {
|
||||
bindAgentSession(agent, runtimeSessionId, chatContext);
|
||||
}
|
||||
@@ -1021,6 +1225,7 @@ public class AgentRunService {
|
||||
runtimeSessionId,
|
||||
chatContext.getUserId() == null ? null : chatContext.getUserId().toString()
|
||||
);
|
||||
Object runtimeEventLock = new Object();
|
||||
AgentRunRegistry.AgentRunContext runContext = new AgentRunRegistry.AgentRunContext(
|
||||
requestId,
|
||||
runtimeSessionId,
|
||||
@@ -1033,15 +1238,40 @@ public class AgentRunService {
|
||||
persistChatlog,
|
||||
owner,
|
||||
lockHandle,
|
||||
event -> handleRuntimeEvent(event, requestId, runOutput, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser, chatContext, finished, persistChatlog),
|
||||
error -> handleRuntimeStreamError(error, requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog),
|
||||
() -> finishRuntimeStream(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)
|
||||
event -> {
|
||||
synchronized (runtimeEventLock) {
|
||||
runRuntimeCallbackSafely(
|
||||
() -> handleRuntimeEvent(event, requestId, runOutput, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser,
|
||||
chatContext, finished, persistChatlog),
|
||||
requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
}
|
||||
},
|
||||
error -> {
|
||||
synchronized (runtimeEventLock) {
|
||||
runRuntimeCallbackSafely(
|
||||
() -> handleRuntimeStreamError(error, requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser,
|
||||
finished, persistChatlog),
|
||||
requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
}
|
||||
},
|
||||
() -> {
|
||||
synchronized (runtimeEventLock) {
|
||||
runRuntimeCallbackSafely(
|
||||
() -> finishRuntimeStream(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser,
|
||||
finished, persistChatlog),
|
||||
requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
}
|
||||
}
|
||||
);
|
||||
agentRunRegistry.register(runContext);
|
||||
lockHandle = null;
|
||||
if (isAguiCancellationRequested(runOutput)) {
|
||||
runContext.eventConsumer().accept(cancellationEvent("用户已停止生成"));
|
||||
return;
|
||||
}
|
||||
if (finished.get()) {
|
||||
runContext.cancel();
|
||||
agentRunRegistry.remove(requestId);
|
||||
@@ -1056,16 +1286,141 @@ public class AgentRunService {
|
||||
} catch (Exception e) {
|
||||
AgentRunRegistry.AgentRunContext runContext = agentRunRegistry.get(requestId);
|
||||
if (runContext != null) {
|
||||
runContext.cancel();
|
||||
try {
|
||||
runContext.cancel();
|
||||
} catch (RuntimeException cancelException) {
|
||||
LOG.warn("Cancel Agent runtime after initialization failure failed, requestId={}",
|
||||
requestId, cancelException);
|
||||
}
|
||||
}
|
||||
try {
|
||||
agentRunRegistry.remove(requestId);
|
||||
} catch (RuntimeException removeException) {
|
||||
LOG.warn("Remove Agent runtime after initialization failure failed, requestId={}",
|
||||
requestId, removeException);
|
||||
}
|
||||
agentRunRegistry.remove(requestId);
|
||||
if (lockHandle != null) {
|
||||
lockHandle.release();
|
||||
releaseRunLockQuietly(lockHandle, requestId);
|
||||
}
|
||||
handleRuntimeError(e, requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 尽力释放会话运行锁，避免 Redis 故障遮蔽随后必须执行的终态与资源清理。
|
||||
*
|
||||
* @param lockHandle 会话锁句柄
|
||||
* @param requestId 运行请求 ID
|
||||
*/
|
||||
private void releaseRunLockQuietly(AgentRunLock.Handle lockHandle, String requestId) {
|
||||
if (lockHandle == null) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
lockHandle.release();
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.warn("Release Agent run lock failed, requestId={}", requestId, exception);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 执行序列化后的 Runtime 回调，并保证终态路径自身抛错时仍释放运行资源。
|
||||
*
|
||||
* @param callback Runtime 事件、错误或完成回调
|
||||
* @param requestId 运行请求 ID
|
||||
* @param runOutput 运行输出
|
||||
* @param chatContext 聊天上下文
|
||||
* @param finished 共享终态标记
|
||||
* @param persistChatlog 是否持久化聊天日志
|
||||
*/
|
||||
private void runRuntimeCallbackSafely(Runnable callback,
|
||||
String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
try {
|
||||
callback.run();
|
||||
} catch (RuntimeException exception) {
|
||||
handleRuntimeCallbackFailure(exception, requestId, runOutput,
|
||||
chatContext, finished, persistChatlog);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 强制收口 Runtime 回调内部故障，不再参与已经完成的终态 CAS。
|
||||
*
|
||||
* @param error 回调内部故障
|
||||
* @param requestId 运行请求 ID
|
||||
* @param runOutput 运行输出
|
||||
* @param chatContext 聊天上下文
|
||||
* @param finished 共享终态标记
|
||||
* @param persistChatlog 是否持久化聊天日志
|
||||
*/
|
||||
private void handleRuntimeCallbackFailure(Throwable error,
|
||||
String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
finished.set(true);
|
||||
Throwable safeError = error == null
|
||||
? new BusinessException("Agent 运行终态处理失败")
|
||||
: error;
|
||||
LOG.error("Agent runtime callback failed, requestId={}, message={}",
|
||||
requestId, safeErrorMessage(safeError), safeError);
|
||||
AgentRunRegistry.AgentRunContext runContext = agentRunRegistry.get(requestId);
|
||||
if (runContext != null) {
|
||||
try {
|
||||
runContext.cancel();
|
||||
} catch (RuntimeException cancelException) {
|
||||
LOG.warn("Cancel Agent runtime after callback failure failed, requestId={}",
|
||||
requestId, cancelException);
|
||||
}
|
||||
}
|
||||
try {
|
||||
agentRunRegistry.remove(requestId);
|
||||
} catch (RuntimeException removeException) {
|
||||
if (runContext != null) {
|
||||
runContext.closeRuntime();
|
||||
}
|
||||
LOG.error("Remove Agent runtime after callback failure failed, requestId={}",
|
||||
requestId, removeException);
|
||||
}
|
||||
try {
|
||||
cancelPending(requestId, safeErrorMessage(safeError), persistChatlog);
|
||||
} catch (RuntimeException pendingException) {
|
||||
LOG.error("Cancel pending approvals after runtime callback failure failed, requestId={}",
|
||||
requestId, pendingException);
|
||||
}
|
||||
if (persistChatlog) {
|
||||
try {
|
||||
chatRuntimeManager.recordFailure(chatContext, safeError);
|
||||
} catch (RuntimeException recordException) {
|
||||
LOG.error("Record Agent runtime callback failure failed, requestId={}",
|
||||
requestId, recordException);
|
||||
}
|
||||
}
|
||||
try {
|
||||
runOutput.completeWithError(safeError);
|
||||
} catch (RuntimeException outputException) {
|
||||
LOG.error("Complete Agent runtime output after callback failure failed, requestId={}",
|
||||
requestId, outputException);
|
||||
}
|
||||
}
|
||||
|
||||
private boolean isAguiCancellationRequested(AgentRunOutput runOutput) {
|
||||
return runOutput instanceof AguiAgentRunOutput aguiOutput
|
||||
&& agentAguiRunStore.isCancellationRequested(aguiOutput.runId());
|
||||
}
|
||||
|
||||
private AgentRuntimeEvent cancellationEvent(String reason) {
|
||||
AgentRuntimeEvent event = AgentRuntimeEvent.of(AgentRuntimeEventType.CANCELLED);
|
||||
event.getPayload().put("reason",
|
||||
reason == null || reason.isBlank() ? "用户已停止生成" : reason);
|
||||
return event;
|
||||
}
|
||||
|
||||
/**
|
||||
* 取消当前节点上指定 Agent 的全部运行。
|
||||
*
|
||||
@@ -1125,6 +1480,9 @@ public class AgentRunService {
|
||||
LegacyThinkingTagParser legacyThinkingTagParser,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
if (!runOutput.cancelRunOnDisconnect()) {
|
||||
return;
|
||||
}
|
||||
Runnable cancelTask = () -> cancelDisconnectedRun(requestId, chatContext, answer,
|
||||
assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
SseEmitter emitter = runOutput.emitter();
|
||||
@@ -1189,6 +1547,15 @@ public class AgentRunService {
|
||||
if (event == null || event.getEventType() == null) {
|
||||
return;
|
||||
}
|
||||
boolean terminalEvent = isTerminalRuntimeEvent(event.getEventType());
|
||||
if (terminalEvent) {
|
||||
// 取消请求来自 HTTP 线程，可能与 Reactor 的完成/失败回调并发；终态必须先统一仲裁再投影。
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
} else if (finished.get()) {
|
||||
return;
|
||||
}
|
||||
Map<String, Object> artifact = buildArtifactPublishedPayload(event);
|
||||
if (artifact != null) {
|
||||
Map<String, Object> persistedPayload = new LinkedHashMap<>();
|
||||
@@ -1373,10 +1740,8 @@ public class AgentRunService {
|
||||
return;
|
||||
}
|
||||
if (event.getEventType() == AgentRuntimeEventType.COMPLETED) {
|
||||
if (!emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)) {
|
||||
return;
|
||||
}
|
||||
emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
String finalText = stringPayload(event, "text");
|
||||
if (!legacyThinkingTagParser.isLegacyFormatDetected() && finalText != null && !finalText.isBlank()) {
|
||||
answer.setLength(0);
|
||||
@@ -1385,52 +1750,47 @@ public class AgentRunService {
|
||||
List<Map<String, Object>> citations = buildKnowledgeCitationPayload(event);
|
||||
if (!citations.isEmpty()) {
|
||||
if (!sendEnvelope(runOutput, ChatDomain.BUSINESS, ChatType.CITATIONS, Map.of("items", citations))) {
|
||||
cancelDisconnectedRun(requestId, chatContext, answer, assistantAccumulator,
|
||||
legacyThinkingTagParser, finished, persistChatlog);
|
||||
return;
|
||||
LOG.debug("Skip completed-run citations after output closed, requestId={}", requestId);
|
||||
}
|
||||
}
|
||||
if (!runOutput.emitRuntimeEvent(event)) {
|
||||
cancelDisconnectedRun(requestId, chatContext, answer, assistantAccumulator,
|
||||
legacyThinkingTagParser, finished, persistChatlog);
|
||||
return;
|
||||
}
|
||||
finishIfNeeded(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, finished, persistChatlog, citations);
|
||||
runOutput.emitRuntimeEvent(event);
|
||||
finishClaimedRun(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, persistChatlog, citations);
|
||||
return;
|
||||
}
|
||||
if (event.getEventType() == AgentRuntimeEventType.CANCELLED) {
|
||||
if (!emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)) {
|
||||
return;
|
||||
}
|
||||
if (!runOutput.emitRuntimeEvent(event)) {
|
||||
cancelDisconnectedRun(requestId, chatContext, answer, assistantAccumulator,
|
||||
legacyThinkingTagParser, finished, persistChatlog);
|
||||
return;
|
||||
}
|
||||
handleRuntimeCancelled(event, requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, finished, persistChatlog);
|
||||
emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
runOutput.emitRuntimeEvent(event);
|
||||
handleClaimedRuntimeCancelled(event, requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, persistChatlog);
|
||||
return;
|
||||
}
|
||||
if (event.getEventType() == AgentRuntimeEventType.FAILED) {
|
||||
if (!emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)) {
|
||||
return;
|
||||
}
|
||||
if (!runOutput.emitRuntimeEvent(event)) {
|
||||
cancelDisconnectedRun(requestId, chatContext, answer, assistantAccumulator,
|
||||
legacyThinkingTagParser, finished, persistChatlog);
|
||||
return;
|
||||
}
|
||||
emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
runOutput.emitRuntimeEvent(event);
|
||||
assistantAccumulator.finalizePendingSkillInvocations("FAILED", "技能调用失败");
|
||||
if (persistChatlog) {
|
||||
recordPartialAssistantIfPresent(chatContext, answer, assistantAccumulator, errorMessage(event));
|
||||
}
|
||||
handleRuntimeError(new BusinessException(errorMessage(event)), requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
handleClaimedRuntimeError(new BusinessException(errorMessage(event)), requestId,
|
||||
runOutput, chatContext, persistChatlog);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断事件是否会收口一次运行。
|
||||
*
|
||||
* @param type 运行时事件类型
|
||||
* @return 完成、取消或失败事件返回 {@code true}
|
||||
*/
|
||||
private boolean isTerminalRuntimeEvent(AgentRuntimeEventType type) {
|
||||
return type == AgentRuntimeEventType.COMPLETED
|
||||
|| type == AgentRuntimeEventType.CANCELLED
|
||||
|| type == AgentRuntimeEventType.FAILED;
|
||||
}
|
||||
|
||||
private boolean deferRuntimeOutput(AgentRuntimeEventType type) {
|
||||
return type == AgentRuntimeEventType.MESSAGE_DELTA
|
||||
|| type == AgentRuntimeEventType.REASONING_STARTED
|
||||
@@ -1584,12 +1944,23 @@ public class AgentRunService {
|
||||
LegacyThinkingTagParser legacyThinkingTagParser,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
if (!emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)) {
|
||||
AgentRunRegistry.AgentRunContext runContext = agentRunRegistry.get(requestId);
|
||||
if (runContext != null && runContext.isSuspended()) {
|
||||
LOG.info("Agent runtime stream suspended, keep SSE and runtime active, requestId={}", requestId);
|
||||
return;
|
||||
}
|
||||
finishIfNeeded(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, finished, persistChatlog);
|
||||
if (!runOutput.canFinishSuccessfully()) {
|
||||
handleRuntimeError(new BusinessException("Agent 事件流缺少完成事件"),
|
||||
requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
return;
|
||||
}
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
finishClaimedRun(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, persistChatlog, List.of());
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -1614,15 +1985,16 @@ public class AgentRunService {
|
||||
LegacyThinkingTagParser legacyThinkingTagParser,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
if (!emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog)) {
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
emitAssistantSegments(legacyThinkingTagParser.finish(), requestId, runOutput, chatContext,
|
||||
answer, assistantAccumulator, legacyThinkingTagParser, finished, persistChatlog);
|
||||
assistantAccumulator.finalizePendingSkillInvocations("FAILED", "技能调用失败");
|
||||
if (persistChatlog) {
|
||||
recordPartialAssistantIfPresent(chatContext, answer, assistantAccumulator, safeErrorMessage(error));
|
||||
}
|
||||
handleRuntimeError(error, requestId, runOutput, chatContext, finished, persistChatlog);
|
||||
handleClaimedRuntimeError(error, requestId, runOutput, chatContext, persistChatlog);
|
||||
}
|
||||
|
||||
private void finishIfNeeded(String requestId,
|
||||
@@ -1657,16 +2029,54 @@ public class AgentRunService {
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
finishClaimedRun(requestId, runOutput, chatContext, answer,
|
||||
assistantAccumulator, persistChatlog, citations);
|
||||
}
|
||||
|
||||
/**
|
||||
* 完成已经赢得终态仲裁的运行。
|
||||
*
|
||||
* @param requestId 运行请求 ID
|
||||
* @param runOutput 运行输出
|
||||
* @param chatContext 聊天上下文
|
||||
* @param answer 最终正文
|
||||
* @param assistantAccumulator 助手结构化累计器
|
||||
* @param persistChatlog 是否持久化聊天记录
|
||||
* @param citations 知识引用
|
||||
*/
|
||||
private void finishClaimedRun(String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
StringBuilder answer,
|
||||
ChatAssistantAccumulator assistantAccumulator,
|
||||
boolean persistChatlog,
|
||||
List<Map<String, Object>> citations) {
|
||||
agentRunRegistry.remove(requestId);
|
||||
String finalAnswer = answer.toString();
|
||||
LOG.info("Agent runtime final answer, requestId={}, sessionId={}, answerLength={}, answer={}",
|
||||
requestId, chatContext.getSessionId(), finalAnswer.length(), toVisibleLogText(finalAnswer));
|
||||
if (persistChatlog) {
|
||||
chatRuntimeManager.recordAssistantCompleted(chatContext,
|
||||
buildAssistantRuntimeMessage(chatContext, finalAnswer, assistantAccumulator, citations));
|
||||
try {
|
||||
chatRuntimeManager.recordAssistantCompleted(chatContext,
|
||||
buildAssistantRuntimeMessage(chatContext, finalAnswer, assistantAccumulator, citations));
|
||||
} catch (RuntimeException exception) {
|
||||
handleClaimedRuntimeError(exception, requestId, runOutput, chatContext, true);
|
||||
return;
|
||||
}
|
||||
}
|
||||
if (!sendDone(runOutput, finalAnswer)) {
|
||||
BusinessException outputFailure = new BusinessException("Agent 可恢复输出终态写入失败");
|
||||
cancelPending(requestId, outputFailure.getMessage(), persistChatlog);
|
||||
LOG.error("Agent run completion output failed, requestId={}, sessionId={}",
|
||||
requestId, chatContext.getSessionId(), outputFailure);
|
||||
if (persistChatlog) {
|
||||
chatRuntimeManager.recordFailure(chatContext, outputFailure);
|
||||
}
|
||||
return;
|
||||
}
|
||||
if (persistChatlog) {
|
||||
chatRuntimeManager.recordCompleted(chatContext);
|
||||
}
|
||||
sendDone(runOutput, finalAnswer);
|
||||
}
|
||||
|
||||
private void handleRuntimeError(Throwable error,
|
||||
@@ -1678,6 +2088,23 @@ public class AgentRunService {
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
handleClaimedRuntimeError(error, requestId, runOutput, chatContext, persistChatlog);
|
||||
}
|
||||
|
||||
/**
|
||||
* 失败收口已经赢得终态仲裁的运行。
|
||||
*
|
||||
* @param error 运行异常
|
||||
* @param requestId 运行请求 ID
|
||||
* @param runOutput 运行输出
|
||||
* @param chatContext 聊天上下文
|
||||
* @param persistChatlog 是否持久化聊天记录
|
||||
*/
|
||||
private void handleClaimedRuntimeError(Throwable error,
|
||||
String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
boolean persistChatlog) {
|
||||
agentRunRegistry.remove(requestId);
|
||||
cancelPending(requestId, safeErrorMessage(error), persistChatlog);
|
||||
Throwable safeError = error == null ? new BusinessException("Agent 运行失败") : error;
|
||||
@@ -1700,17 +2127,24 @@ public class AgentRunService {
|
||||
return error.getMessage();
|
||||
}
|
||||
|
||||
private void handleRuntimeCancelled(AgentRuntimeEvent event,
|
||||
String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
StringBuilder answer,
|
||||
ChatAssistantAccumulator assistantAccumulator,
|
||||
AtomicBoolean finished,
|
||||
boolean persistChatlog) {
|
||||
if (!finished.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
/**
|
||||
* 取消收口已经赢得终态仲裁的运行。
|
||||
*
|
||||
* @param event 取消事件
|
||||
* @param requestId 运行请求 ID
|
||||
* @param runOutput 运行输出
|
||||
* @param chatContext 聊天上下文
|
||||
* @param answer 已生成正文
|
||||
* @param assistantAccumulator 助手结构化累计器
|
||||
* @param persistChatlog 是否持久化聊天记录
|
||||
*/
|
||||
private void handleClaimedRuntimeCancelled(AgentRuntimeEvent event,
|
||||
String requestId,
|
||||
AgentRunOutput runOutput,
|
||||
ChatRuntimeContext chatContext,
|
||||
StringBuilder answer,
|
||||
ChatAssistantAccumulator assistantAccumulator,
|
||||
boolean persistChatlog) {
|
||||
agentRunRegistry.remove(requestId);
|
||||
String reason = errorMessage(event);
|
||||
assistantAccumulator.finalizePendingSkillInvocations("CANCELLED", "技能调用已停止");
|
||||
|
||||
@@ -0,0 +1,32 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
/**
|
||||
* AG-UI 可重连运行的公开状态与内部路由描述。
|
||||
*
|
||||
* @param runId 客户端运行 ID
|
||||
* @param requestId 服务端内部请求 ID
|
||||
* @param threadId AG-UI thread ID
|
||||
* @param agentId Agent ID
|
||||
* @param sessionId 运行会话 ID
|
||||
* @param userId 所属用户 ID
|
||||
* @param tenantId 所属租户 ID
|
||||
* @param draft 是否为草稿试运行
|
||||
* @param status 当前运行状态
|
||||
* @param lastCursor 已持久化的最后事件游标
|
||||
* @param createdAt 创建时间戳
|
||||
* @param updatedAt 最近更新时间戳
|
||||
*/
|
||||
public record AgentAguiRunDescriptor(
|
||||
String runId,
|
||||
String requestId,
|
||||
String threadId,
|
||||
String agentId,
|
||||
String sessionId,
|
||||
String userId,
|
||||
String tenantId,
|
||||
boolean draft,
|
||||
AgentAguiRunStatus status,
|
||||
long lastCursor,
|
||||
long createdAt,
|
||||
long updatedAt) {
|
||||
}
|
||||
@@ -0,0 +1,330 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import com.fasterxml.jackson.databind.JsonNode;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import com.fasterxml.jackson.databind.node.ObjectNode;
|
||||
import com.easyagents.agui.AguiExtendedEvent;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.scheduling.annotation.Scheduled;
|
||||
import org.springframework.stereotype.Component;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.atomic.AtomicLong;
|
||||
import java.util.concurrent.locks.ReentrantLock;
|
||||
import java.util.UUID;
|
||||
|
||||
/**
|
||||
* 将高频 AG-UI delta 在内存中短暂合并并批量写入 Redis。
|
||||
*
|
||||
* <p>该组件以每个 runId 的同步缓冲保证事件顺序；Redis 故障时保留原批次等待下一次重试，
|
||||
* 避免每个模型 token 都产生一次网络往返。</p>
|
||||
*/
|
||||
@Component
|
||||
public class AgentAguiRunJournal {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(AgentAguiRunJournal.class);
|
||||
private static final int MAX_BUFFERED_EVENTS_PER_RUN = 4096;
|
||||
private static final long MAX_EVENT_BYTES = 512L * 1024L;
|
||||
private static final long MAX_BUFFERED_BYTES_PER_RUN = 4L * 1024L * 1024L;
|
||||
private static final long MAX_BUFFERED_BYTES_GLOBAL = 64L * 1024L * 1024L;
|
||||
private static final long COMPLETED_RETRY_WINDOW_MILLIS = 5L * 60L * 1000L;
|
||||
|
||||
private final AgentAguiRunStore store;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final Map<String, RunBuffer> buffers = new ConcurrentHashMap<>();
|
||||
private final AtomicLong globalBufferedBytes = new AtomicLong();
|
||||
|
||||
/**
|
||||
* 创建 AG-UI 事件日志协调器。
|
||||
*
|
||||
* @param store Redis 运行存储
|
||||
* @param objectMapper JSON 映射器
|
||||
*/
|
||||
public AgentAguiRunJournal(AgentAguiRunStore store, ObjectMapper objectMapper) {
|
||||
this.store = store;
|
||||
this.objectMapper = objectMapper;
|
||||
}
|
||||
|
||||
/**
|
||||
* 追加一条已编码协议事件。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param data 协议事件 JSON
|
||||
*/
|
||||
public void append(String runId, String data) {
|
||||
if (data == null || data.isBlank()) {
|
||||
throw new IllegalArgumentException("AG-UI 事件不能为空");
|
||||
}
|
||||
long dataBytes = utf8Bytes(data);
|
||||
if (dataBytes > MAX_EVENT_BYTES) {
|
||||
throw new IllegalStateException("单条 AG-UI 事件超过 512 KiB 安全上限");
|
||||
}
|
||||
RunBuffer buffer = buffers.computeIfAbsent(runId, ignored -> new RunBuffer());
|
||||
try {
|
||||
synchronized (buffer) {
|
||||
if (buffer.completed) {
|
||||
throw new IllegalStateException("AG-UI 运行事件日志已经收口");
|
||||
}
|
||||
int pendingEventCount = buffer.pendingBatch == null ? 0 : buffer.pendingBatch.events().size();
|
||||
if (pendingEventCount + buffer.tailEvents.size() >= MAX_BUFFERED_EVENTS_PER_RUN) {
|
||||
throw new IllegalStateException("AG-UI 事件日志暂时不可用，请稍后重试");
|
||||
}
|
||||
String merged = buffer.tailEvents.isEmpty()
|
||||
? null
|
||||
: mergeDelta(buffer.tailEvents.get(buffer.tailEvents.size() - 1), data);
|
||||
long previousBytes = 0L;
|
||||
long nextBytes = dataBytes;
|
||||
if (merged != null) {
|
||||
String previous = buffer.tailEvents.get(buffer.tailEvents.size() - 1);
|
||||
previousBytes = utf8Bytes(previous);
|
||||
nextBytes = utf8Bytes(merged);
|
||||
if (nextBytes > MAX_EVENT_BYTES) {
|
||||
throw new IllegalStateException("合并后的 AG-UI 事件超过 512 KiB 安全上限");
|
||||
}
|
||||
}
|
||||
long deltaBytes = nextBytes - previousBytes;
|
||||
if (buffer.totalBytes() + deltaBytes > MAX_BUFFERED_BYTES_PER_RUN) {
|
||||
throw new IllegalStateException("单次 AG-UI 运行待写日志超过 4 MiB 安全上限");
|
||||
}
|
||||
long globalBytes = globalBufferedBytes.addAndGet(deltaBytes);
|
||||
if (globalBytes > MAX_BUFFERED_BYTES_GLOBAL) {
|
||||
globalBufferedBytes.addAndGet(-deltaBytes);
|
||||
throw new IllegalStateException("AG-UI 待写日志超过 64 MiB 全局安全上限");
|
||||
}
|
||||
if (merged == null) {
|
||||
buffer.tailEvents.add(data);
|
||||
} else {
|
||||
buffer.tailEvents.set(buffer.tailEvents.size() - 1, merged);
|
||||
}
|
||||
buffer.tailBytes += deltaBytes;
|
||||
}
|
||||
} catch (RuntimeException exception) {
|
||||
removeEmptyBuffer(runId, buffer);
|
||||
throw exception;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 丢弃尚未确认持久化的普通事件，并以独立小批次记录日志故障终态。
|
||||
*
|
||||
* <p>该路径保留已经成功写入 Redis 的历史。即使前一批次处于结果不确定状态，
|
||||
* 新终态也使用新的批次 ID，避免幂等判重吞掉故障事件。</p>
|
||||
*
|
||||
* @param threadId AG-UI 线程 ID
|
||||
* @param runId 运行 ID
|
||||
* @param code 稳定错误码
|
||||
* @param message 用户可理解的错误信息
|
||||
*/
|
||||
public void fail(String threadId, String runId, String code, String message) {
|
||||
String terminalEvent;
|
||||
try {
|
||||
terminalEvent = objectMapper.writeValueAsString(new AguiExtendedEvent.RunError(
|
||||
threadId, runId, message, code));
|
||||
} catch (Exception exception) {
|
||||
throw new IllegalStateException("AG-UI 日志故障终态编码失败", exception);
|
||||
}
|
||||
long terminalBytes = utf8Bytes(terminalEvent);
|
||||
RunBuffer buffer = buffers.computeIfAbsent(runId, ignored -> new RunBuffer());
|
||||
synchronized (buffer) {
|
||||
long discardedBytes = buffer.totalBytes();
|
||||
buffer.pendingBatch = null;
|
||||
buffer.tailEvents.clear();
|
||||
buffer.tailBytes = 0L;
|
||||
if (discardedBytes > 0L) {
|
||||
globalBufferedBytes.addAndGet(-discardedBytes);
|
||||
}
|
||||
buffer.pendingBatch = new PendingBatch(
|
||||
UUID.randomUUID().toString(), List.of(terminalEvent), terminalBytes);
|
||||
globalBufferedBytes.addAndGet(terminalBytes);
|
||||
buffer.completed = true;
|
||||
buffer.completedAt = System.currentTimeMillis();
|
||||
}
|
||||
flush(runId, buffer, true);
|
||||
}
|
||||
|
||||
/**
|
||||
* 在运行收口前同步尝试刷入该运行的剩余事件。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
*/
|
||||
public void flush(String runId) {
|
||||
RunBuffer buffer = buffers.get(runId);
|
||||
if (buffer != null) {
|
||||
flush(runId, buffer, true);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 同步刷入终态事件，并在确认缓冲为空后释放本地运行缓冲。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
*/
|
||||
public void complete(String runId) {
|
||||
RunBuffer buffer = buffers.get(runId);
|
||||
if (buffer == null) {
|
||||
return;
|
||||
}
|
||||
synchronized (buffer) {
|
||||
buffer.completed = true;
|
||||
buffer.completedAt = System.currentTimeMillis();
|
||||
}
|
||||
flush(runId, buffer, true);
|
||||
}
|
||||
|
||||
/**
|
||||
* 定期批量写入所有活动运行的待持久化事件。
|
||||
*/
|
||||
@Scheduled(fixedDelayString = "${easyflow.agent.runtime.agui-journal-flush-interval:50ms}")
|
||||
public void flushPending() {
|
||||
for (Map.Entry<String, RunBuffer> entry : buffers.entrySet()) {
|
||||
flush(entry.getKey(), entry.getValue(), false);
|
||||
}
|
||||
}
|
||||
|
||||
private void flush(String runId, RunBuffer buffer, boolean propagateFailure) {
|
||||
boolean acquired;
|
||||
if (propagateFailure) {
|
||||
buffer.flushLock.lock();
|
||||
acquired = true;
|
||||
} else {
|
||||
acquired = buffer.flushLock.tryLock();
|
||||
}
|
||||
if (!acquired) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
while (true) {
|
||||
PendingBatch batch;
|
||||
synchronized (buffer) {
|
||||
if (buffer.pendingBatch == null && !buffer.tailEvents.isEmpty()) {
|
||||
buffer.pendingBatch = new PendingBatch(
|
||||
UUID.randomUUID().toString(),
|
||||
List.copyOf(buffer.tailEvents),
|
||||
buffer.tailBytes);
|
||||
buffer.tailEvents.clear();
|
||||
buffer.tailBytes = 0L;
|
||||
}
|
||||
batch = buffer.pendingBatch;
|
||||
if (batch == null) {
|
||||
if (buffer.completed) {
|
||||
removeBuffer(runId, buffer);
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
try {
|
||||
// Redis I/O 必须位于 RunBuffer 锁外，避免阻塞模型事件线程。
|
||||
store.append(runId, batch.id(), batch.events());
|
||||
synchronized (buffer) {
|
||||
if (buffer.pendingBatch == batch) {
|
||||
buffer.pendingBatch = null;
|
||||
globalBufferedBytes.addAndGet(-batch.bytes());
|
||||
}
|
||||
}
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("AG-UI 事件日志写入失败: runId={}, eventCount={}",
|
||||
runId, batch.events().size(), exception);
|
||||
boolean batchWasReplaced;
|
||||
synchronized (buffer) {
|
||||
batchWasReplaced = buffer.pendingBatch != batch;
|
||||
if (!batchWasReplaced && buffer.completed && buffer.completedAt > 0L
|
||||
&& System.currentTimeMillis() - buffer.completedAt
|
||||
>= COMPLETED_RETRY_WINDOW_MILLIS) {
|
||||
LOG.error("丢弃超过重试窗口的 AG-UI 终态日志缓冲: runId={}, eventCount={}, bytes={}",
|
||||
runId, buffer.eventCount(), buffer.totalBytes());
|
||||
removeBuffer(runId, buffer);
|
||||
}
|
||||
}
|
||||
if (batchWasReplaced) {
|
||||
continue;
|
||||
}
|
||||
if (propagateFailure) {
|
||||
throw exception;
|
||||
}
|
||||
return;
|
||||
}
|
||||
}
|
||||
} finally {
|
||||
buffer.flushLock.unlock();
|
||||
}
|
||||
}
|
||||
|
||||
private void removeEmptyBuffer(String runId, RunBuffer buffer) {
|
||||
synchronized (buffer) {
|
||||
if (buffer.pendingBatch == null && buffer.tailEvents.isEmpty() && !buffer.flushLock.isLocked()) {
|
||||
buffers.remove(runId, buffer);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private void removeBuffer(String runId, RunBuffer buffer) {
|
||||
if (!buffers.remove(runId, buffer)) {
|
||||
return;
|
||||
}
|
||||
long remainingBytes = buffer.totalBytes();
|
||||
buffer.pendingBatch = null;
|
||||
buffer.tailEvents.clear();
|
||||
buffer.tailBytes = 0L;
|
||||
if (remainingBytes > 0L) {
|
||||
globalBufferedBytes.addAndGet(-remainingBytes);
|
||||
}
|
||||
}
|
||||
|
||||
private long utf8Bytes(String value) {
|
||||
return value.getBytes(StandardCharsets.UTF_8).length;
|
||||
}
|
||||
|
||||
private String mergeDelta(String previous, String current) {
|
||||
try {
|
||||
JsonNode previousNode = objectMapper.readTree(previous);
|
||||
JsonNode currentNode = objectMapper.readTree(current);
|
||||
String type = currentNode.path("type").asText();
|
||||
if (!type.equals(previousNode.path("type").asText()) || !isMergeable(type)) {
|
||||
return null;
|
||||
}
|
||||
String identityField = "TOOL_CALL_ARGS".equals(type) ? "toolCallId" : "messageId";
|
||||
if (!currentNode.path(identityField).asText().equals(previousNode.path(identityField).asText())) {
|
||||
return null;
|
||||
}
|
||||
if (!(previousNode instanceof ObjectNode previousObject)) {
|
||||
return null;
|
||||
}
|
||||
previousObject.put("delta",
|
||||
previousNode.path("delta").asText() + currentNode.path("delta").asText());
|
||||
return objectMapper.writeValueAsString(previousObject);
|
||||
} catch (Exception ignored) {
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
private boolean isMergeable(String type) {
|
||||
return "TEXT_MESSAGE_CONTENT".equals(type)
|
||||
|| "REASONING_MESSAGE_CONTENT".equals(type)
|
||||
|| "TOOL_CALL_ARGS".equals(type);
|
||||
}
|
||||
|
||||
private record PendingBatch(String id, List<String> events, long bytes) {
|
||||
}
|
||||
|
||||
private static final class RunBuffer {
|
||||
private final List<String> tailEvents = new ArrayList<>();
|
||||
private final ReentrantLock flushLock = new ReentrantLock();
|
||||
private PendingBatch pendingBatch;
|
||||
private long tailBytes;
|
||||
private boolean completed;
|
||||
private long completedAt;
|
||||
|
||||
private int eventCount() {
|
||||
return tailEvents.size() + (pendingBatch == null ? 0 : pendingBatch.events().size());
|
||||
}
|
||||
|
||||
private long totalBytes() {
|
||||
return tailBytes + (pendingBatch == null ? 0L : pendingBatch.bytes());
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
/**
|
||||
* AG-UI 可重连运行状态。
|
||||
*/
|
||||
public enum AgentAguiRunStatus {
|
||||
|
||||
/**
|
||||
* 运行正在生成或等待人工审批。
|
||||
*/
|
||||
RUNNING,
|
||||
|
||||
/**
|
||||
* 运行成功完成。
|
||||
*/
|
||||
COMPLETED,
|
||||
|
||||
/**
|
||||
* 运行失败。
|
||||
*/
|
||||
FAILED,
|
||||
|
||||
/**
|
||||
* 运行被用户显式取消。
|
||||
*/
|
||||
CANCELLED;
|
||||
|
||||
/**
|
||||
* 判断当前状态是否已经终止。
|
||||
*
|
||||
* @return 已终止时为 true
|
||||
*/
|
||||
public boolean isTerminal() {
|
||||
return this != RUNNING;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,36 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
/**
|
||||
* 提供给前端刷新恢复判断的 AG-UI 运行状态。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param threadId AG-UI thread ID
|
||||
* @param status 当前状态
|
||||
* @param lastCursor 最后事件游标
|
||||
* @param draft 是否为草稿试运行
|
||||
* @param updatedAt 最近更新时间戳
|
||||
*/
|
||||
public record AgentAguiRunStatusView(
|
||||
String runId,
|
||||
String threadId,
|
||||
AgentAguiRunStatus status,
|
||||
long lastCursor,
|
||||
boolean draft,
|
||||
long updatedAt) {
|
||||
|
||||
/**
|
||||
* 从内部运行描述生成公开状态。
|
||||
*
|
||||
* @param descriptor 内部运行描述
|
||||
* @return 公开状态
|
||||
*/
|
||||
public static AgentAguiRunStatusView from(AgentAguiRunDescriptor descriptor) {
|
||||
return new AgentAguiRunStatusView(
|
||||
descriptor.runId(),
|
||||
descriptor.threadId(),
|
||||
descriptor.status(),
|
||||
descriptor.lastCursor(),
|
||||
descriptor.draft(),
|
||||
descriptor.updatedAt());
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,369 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import com.fasterxml.jackson.databind.JsonNode;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.springframework.data.redis.core.HashOperations;
|
||||
import org.springframework.data.redis.core.StringRedisTemplate;
|
||||
import org.springframework.data.redis.core.script.DefaultRedisScript;
|
||||
import org.springframework.stereotype.Component;
|
||||
import com.easyagents.agui.AguiExtendedEvent;
|
||||
import tech.easyflow.agent.config.AgentRuntimeProperties;
|
||||
import tech.easyflow.common.entity.LoginAccount;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.UUID;
|
||||
import java.util.regex.Pattern;
|
||||
|
||||
/**
|
||||
* 基于 Redis 的 AG-UI 运行元数据与顺序事件日志。
|
||||
*
|
||||
* <p>运行 ID 只承担幂等查找作用，所有读取和取消仍必须校验登录用户与租户归属。</p>
|
||||
*/
|
||||
@Component
|
||||
public class AgentAguiRunStore {
|
||||
|
||||
private static final String RUN_PREFIX = "easyflow:agent:agui:run:";
|
||||
private static final String STATUS_FIELD = "status";
|
||||
private static final int REPLAY_BATCH_SIZE = 256;
|
||||
private static final Pattern RUN_ID_PATTERN = Pattern.compile("[A-Za-z0-9_-]{1,128}");
|
||||
private static final long MAX_RUN_EVENT_BYTES = 16L * 1024L * 1024L;
|
||||
private static final long TERMINAL_EVENT_RESERVE_BYTES = 64L * 1024L;
|
||||
private static final DefaultRedisScript<Long> APPEND_SCRIPT = new DefaultRedisScript<>("""
|
||||
local status = redis.call('HGET', KEYS[1], 'status')
|
||||
if not status then
|
||||
return -3
|
||||
end
|
||||
local lastCursor = tonumber(redis.call('HGET', KEYS[1], 'lastCursor') or '0')
|
||||
if redis.call('HGET', KEYS[1], 'lastBatchId') == ARGV[1] then
|
||||
return lastCursor
|
||||
end
|
||||
if status ~= 'RUNNING' then
|
||||
return lastCursor
|
||||
end
|
||||
local currentBytes = tonumber(redis.call('HGET', KEYS[1], 'eventBytes') or '0')
|
||||
local incomingBytes = tonumber(ARGV[5])
|
||||
local byteLimit = tonumber(ARGV[6])
|
||||
if ARGV[4] ~= '' then
|
||||
byteLimit = byteLimit + tonumber(ARGV[7])
|
||||
end
|
||||
if currentBytes + incomingBytes > byteLimit then
|
||||
return -1
|
||||
end
|
||||
local eventCount = tonumber(ARGV[8])
|
||||
for index = 1, eventCount do
|
||||
redis.call('RPUSH', KEYS[2], ARGV[8 + index])
|
||||
end
|
||||
local cursor = redis.call('LLEN', KEYS[2])
|
||||
redis.call('HSET', KEYS[1],
|
||||
'updatedAt', ARGV[2],
|
||||
'lastBatchId', ARGV[1],
|
||||
'eventBytes', currentBytes + incomingBytes,
|
||||
'lastCursor', cursor)
|
||||
if ARGV[4] ~= '' then
|
||||
redis.call('HSET', KEYS[1], 'status', ARGV[4])
|
||||
end
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[3])
|
||||
redis.call('PEXPIRE', KEYS[2], ARGV[3])
|
||||
redis.call('PEXPIRE', KEYS[3], ARGV[3])
|
||||
return cursor
|
||||
""", Long.class);
|
||||
private static final DefaultRedisScript<Long> REQUEST_CANCEL_SCRIPT = new DefaultRedisScript<>("""
|
||||
if redis.call('HGET', KEYS[1], 'status') ~= 'RUNNING' then
|
||||
return 0
|
||||
end
|
||||
redis.call('HSET', KEYS[1], 'cancelRequested', 'true', 'updatedAt', ARGV[1])
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[2])
|
||||
return 1
|
||||
""", Long.class);
|
||||
|
||||
private final StringRedisTemplate redisTemplate;
|
||||
private final ObjectMapper objectMapper;
|
||||
private final Duration retention;
|
||||
|
||||
/**
|
||||
* 创建 AG-UI 运行存储。
|
||||
*
|
||||
* @param redisTemplate Redis 字符串模板
|
||||
* @param objectMapper JSON 映射器
|
||||
* @param properties Agent 运行配置
|
||||
*/
|
||||
public AgentAguiRunStore(StringRedisTemplate redisTemplate,
|
||||
ObjectMapper objectMapper,
|
||||
AgentRuntimeProperties properties) {
|
||||
this.redisTemplate = redisTemplate;
|
||||
this.objectMapper = objectMapper;
|
||||
this.retention = properties.getAguiRunRetention();
|
||||
}
|
||||
|
||||
/**
|
||||
* 在模型运行启动前注册唯一运行元数据。
|
||||
*
|
||||
* @param descriptor 初始运行描述
|
||||
* @throws BusinessException runId 已存在时抛出
|
||||
*/
|
||||
public void create(AgentAguiRunDescriptor descriptor) {
|
||||
String runId = requireRunId(descriptor == null ? null : descriptor.runId());
|
||||
Boolean reserved = redisTemplate.opsForValue().setIfAbsent(
|
||||
reservationKey(runId), descriptor.userId(), retention);
|
||||
if (!Boolean.TRUE.equals(reserved)) {
|
||||
throw new BusinessException("当前 Agent 运行标识已存在，请勿重复提交");
|
||||
}
|
||||
try {
|
||||
Map<String, String> values = new LinkedHashMap<>();
|
||||
values.put("runId", runId);
|
||||
values.put("requestId", value(descriptor.requestId()));
|
||||
values.put("threadId", value(descriptor.threadId()));
|
||||
values.put("agentId", value(descriptor.agentId()));
|
||||
values.put("sessionId", value(descriptor.sessionId()));
|
||||
values.put("userId", value(descriptor.userId()));
|
||||
values.put("tenantId", value(descriptor.tenantId()));
|
||||
values.put("draft", Boolean.toString(descriptor.draft()));
|
||||
values.put(STATUS_FIELD, AgentAguiRunStatus.RUNNING.name());
|
||||
values.put("cancelRequested", Boolean.FALSE.toString());
|
||||
values.put("eventBytes", "0");
|
||||
values.put("lastBatchId", "");
|
||||
values.put("lastCursor", "0");
|
||||
values.put("createdAt", Long.toString(descriptor.createdAt()));
|
||||
values.put("updatedAt", Long.toString(descriptor.updatedAt()));
|
||||
meta().putAll(metaKey(runId), values);
|
||||
redisTemplate.expire(metaKey(runId), retention);
|
||||
redisTemplate.expire(eventsKey(runId), retention);
|
||||
} catch (RuntimeException exception) {
|
||||
redisTemplate.delete(reservationKey(runId));
|
||||
throw exception;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 批量追加严格有序的协议事件。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param events 已编码 AG-UI 事件
|
||||
* @return 追加后的最后事件游标
|
||||
*/
|
||||
public long append(String runId, String batchId, List<String> events) {
|
||||
requireRunId(runId);
|
||||
if (events == null || events.isEmpty()) {
|
||||
AgentAguiRunDescriptor descriptor = find(runId);
|
||||
return descriptor == null ? 0L : descriptor.lastCursor();
|
||||
}
|
||||
if (batchId == null || batchId.isBlank()) {
|
||||
throw new IllegalArgumentException("AG-UI 事件批次 ID 不能为空");
|
||||
}
|
||||
AgentAguiRunStatus terminalStatus = terminalStatus(events);
|
||||
long eventBytes = events.stream()
|
||||
.mapToLong(event -> event.getBytes(StandardCharsets.UTF_8).length)
|
||||
.sum();
|
||||
List<String> arguments = new java.util.ArrayList<>(8 + events.size());
|
||||
arguments.add(batchId);
|
||||
arguments.add(Long.toString(System.currentTimeMillis()));
|
||||
arguments.add(Long.toString(retention.toMillis()));
|
||||
arguments.add(terminalStatus == null ? "" : terminalStatus.name());
|
||||
arguments.add(Long.toString(eventBytes));
|
||||
arguments.add(Long.toString(MAX_RUN_EVENT_BYTES));
|
||||
arguments.add(Long.toString(TERMINAL_EVENT_RESERVE_BYTES));
|
||||
arguments.add(Integer.toString(events.size()));
|
||||
arguments.addAll(events);
|
||||
Long cursor = redisTemplate.execute(
|
||||
APPEND_SCRIPT,
|
||||
List.of(metaKey(runId), eventsKey(runId), reservationKey(runId)),
|
||||
arguments.toArray());
|
||||
if (cursor == null || cursor == -3L) {
|
||||
throw new IllegalStateException("AG-UI 运行记录不存在或已过期");
|
||||
}
|
||||
if (cursor == -1L) {
|
||||
throw new IllegalStateException("AG-UI 运行事件超过 16 MiB 安全上限");
|
||||
}
|
||||
return cursor;
|
||||
}
|
||||
|
||||
/**
|
||||
* 读取游标之后的一批事件。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param afterCursor 已消费的最后游标，零表示从头读取
|
||||
* @return 按游标升序排列的事件
|
||||
*/
|
||||
public List<String> readAfter(String runId, long afterCursor) {
|
||||
requireRunId(runId);
|
||||
long start = Math.max(0L, afterCursor);
|
||||
List<String> values = redisTemplate.opsForList().range(
|
||||
eventsKey(runId), start, start + REPLAY_BATCH_SIZE - 1L);
|
||||
return values == null ? List.of() : values;
|
||||
}
|
||||
|
||||
/**
|
||||
* 查询运行描述。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @return 运行描述；不存在时为 null
|
||||
*/
|
||||
public AgentAguiRunDescriptor find(String runId) {
|
||||
requireRunId(runId);
|
||||
Map<Object, Object> values = meta().entries(metaKey(runId));
|
||||
if (values == null || values.isEmpty()) {
|
||||
return null;
|
||||
}
|
||||
return new AgentAguiRunDescriptor(
|
||||
text(values, "runId"),
|
||||
text(values, "requestId"),
|
||||
text(values, "threadId"),
|
||||
text(values, "agentId"),
|
||||
text(values, "sessionId"),
|
||||
text(values, "userId"),
|
||||
text(values, "tenantId"),
|
||||
Boolean.parseBoolean(text(values, "draft")),
|
||||
parseStatus(text(values, STATUS_FIELD)),
|
||||
number(values, "lastCursor"),
|
||||
number(values, "createdAt"),
|
||||
number(values, "updatedAt"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 查询并校验当前登录账号对运行的访问权。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param account 当前登录账号
|
||||
* @return 已授权运行描述
|
||||
* @throws BusinessException 运行不存在或无访问权时抛出
|
||||
*/
|
||||
public AgentAguiRunDescriptor requireOwned(String runId, LoginAccount account) {
|
||||
AgentAguiRunDescriptor descriptor = find(runId);
|
||||
if (descriptor == null) {
|
||||
throw new BusinessException("Agent 运行记录不存在或已过期");
|
||||
}
|
||||
String userId = account == null || account.getId() == null ? null : account.getId().toString();
|
||||
String tenantId = account == null || account.getTenantId() == null ? null : account.getTenantId().toString();
|
||||
if (!descriptor.userId().equals(userId) || !descriptor.tenantId().equals(tenantId)) {
|
||||
throw new BusinessException("无权访问该 Agent 运行");
|
||||
}
|
||||
return descriptor;
|
||||
}
|
||||
|
||||
/**
|
||||
* 将无法继续的运行收口为失败状态。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
*/
|
||||
public void failOwnerLost(AgentAguiRunDescriptor descriptor) {
|
||||
if (descriptor == null || descriptor.status().isTerminal()) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
String event = objectMapper.writeValueAsString(new AguiExtendedEvent.RunError(
|
||||
descriptor.threadId(),
|
||||
descriptor.runId(),
|
||||
"Agent 运行节点已不可用",
|
||||
"AGENT_RUN_OWNER_LOST"));
|
||||
append(descriptor.runId(), "owner-lost-" + UUID.randomUUID(), List.of(event));
|
||||
} catch (RuntimeException exception) {
|
||||
throw exception;
|
||||
} catch (Exception exception) {
|
||||
throw new IllegalStateException("AG-UI owner-lost 终态编码失败", exception);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 记录显式取消意图，供 Runtime 注册前后的启动窗口检查。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @return 本次是否首次接受取消意图
|
||||
*/
|
||||
public boolean requestCancellation(String runId) {
|
||||
requireRunId(runId);
|
||||
Long result = redisTemplate.execute(
|
||||
REQUEST_CANCEL_SCRIPT,
|
||||
List.of(metaKey(runId)),
|
||||
Long.toString(System.currentTimeMillis()),
|
||||
Long.toString(retention.toMillis()));
|
||||
return Long.valueOf(1L).equals(result);
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断运行是否收到显式取消意图。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @return 已请求取消时为 true
|
||||
*/
|
||||
public boolean isCancellationRequested(String runId) {
|
||||
Object value = meta().get(metaKey(requireRunId(runId)), "cancelRequested");
|
||||
return Boolean.parseBoolean(value == null ? "false" : value.toString());
|
||||
}
|
||||
|
||||
private HashOperations<String, Object, Object> meta() {
|
||||
return redisTemplate.opsForHash();
|
||||
}
|
||||
|
||||
private AgentAguiRunStatus terminalStatus(List<String> events) {
|
||||
AgentAguiRunStatus result = null;
|
||||
for (String event : events) {
|
||||
try {
|
||||
JsonNode value = objectMapper.readTree(event);
|
||||
String type = value.path("type").asText();
|
||||
if ("RUN_FINISHED".equals(type)) {
|
||||
result = AgentAguiRunStatus.COMPLETED;
|
||||
} else if ("RUN_ERROR".equals(type)) {
|
||||
result = "RUN_CANCELLED".equals(value.path("code").asText())
|
||||
? AgentAguiRunStatus.CANCELLED
|
||||
: AgentAguiRunStatus.FAILED;
|
||||
}
|
||||
} catch (Exception exception) {
|
||||
throw new IllegalStateException("AG-UI 运行事件格式不合法", exception);
|
||||
}
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
private AgentAguiRunStatus parseStatus(String value) {
|
||||
try {
|
||||
return AgentAguiRunStatus.valueOf(value);
|
||||
} catch (Exception ignored) {
|
||||
return AgentAguiRunStatus.FAILED;
|
||||
}
|
||||
}
|
||||
|
||||
private String requireRunId(String runId) {
|
||||
if (runId == null || !RUN_ID_PATTERN.matcher(runId).matches()) {
|
||||
throw new BusinessException("Agent 运行 ID 不合法");
|
||||
}
|
||||
return runId;
|
||||
}
|
||||
|
||||
private String reservationKey(String runId) {
|
||||
return runKey(runId, "reserved");
|
||||
}
|
||||
|
||||
private String metaKey(String runId) {
|
||||
return runKey(runId, "meta");
|
||||
}
|
||||
|
||||
private String eventsKey(String runId) {
|
||||
return runKey(runId, "events");
|
||||
}
|
||||
|
||||
private String runKey(String runId, String suffix) {
|
||||
return RUN_PREFIX + "{" + runId + "}:" + suffix;
|
||||
}
|
||||
|
||||
private String value(String value) {
|
||||
return value == null ? "" : value;
|
||||
}
|
||||
|
||||
private String text(Map<Object, Object> values, String key) {
|
||||
Object value = values.get(key);
|
||||
return value == null ? "" : value.toString();
|
||||
}
|
||||
|
||||
private long number(Map<Object, Object> values, String key) {
|
||||
try {
|
||||
return Long.parseLong(text(values, key));
|
||||
} catch (NumberFormatException ignored) {
|
||||
return 0L;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,160 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.scheduling.annotation.Scheduled;
|
||||
import org.springframework.stereotype.Component;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
|
||||
import java.io.IOException;
|
||||
import java.util.LinkedHashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.ConcurrentHashMap;
|
||||
import java.util.concurrent.CopyOnWriteArrayList;
|
||||
|
||||
/**
|
||||
* 以共享调度器为本节点上的 AG-UI 重连订阅者增量重放 Redis 事件。
|
||||
*
|
||||
* <p>所有订阅共享一个短周期任务，不为每个浏览器连接创建阻塞线程。</p>
|
||||
*/
|
||||
@Component
|
||||
public class AgentAguiRunSubscriptionService {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(AgentAguiRunSubscriptionService.class);
|
||||
private static final long KEEP_ALIVE_INTERVAL_MILLIS = 15_000L;
|
||||
|
||||
private final AgentAguiRunStore store;
|
||||
private final Map<String, CopyOnWriteArrayList<Subscriber>> subscribers = new ConcurrentHashMap<>();
|
||||
|
||||
/**
|
||||
* 创建 AG-UI 重连订阅服务。
|
||||
*
|
||||
* @param store Redis 运行存储
|
||||
*/
|
||||
public AgentAguiRunSubscriptionService(AgentAguiRunStore store) {
|
||||
this.store = store;
|
||||
}
|
||||
|
||||
/**
|
||||
* 注册一个从指定游标继续的 SSE 订阅。
|
||||
*
|
||||
* @param descriptor 已授权运行描述
|
||||
* @param afterCursor 已消费的最后游标
|
||||
* @return 重连 SSE 发射器
|
||||
*/
|
||||
public SseEmitter subscribe(AgentAguiRunDescriptor descriptor, long afterCursor) {
|
||||
return subscribe(descriptor, afterCursor, new SseEmitter(0L));
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用指定发射器注册订阅，供生命周期测试复用同一调度逻辑。
|
||||
*
|
||||
* @param descriptor 已授权运行描述
|
||||
* @param afterCursor 已消费的最后游标
|
||||
* @param emitter SSE 发射器
|
||||
* @return 已注册发射器
|
||||
*/
|
||||
SseEmitter subscribe(AgentAguiRunDescriptor descriptor,
|
||||
long afterCursor,
|
||||
SseEmitter emitter) {
|
||||
Subscriber subscriber = new Subscriber(emitter, Math.max(0L, afterCursor));
|
||||
subscribers.computeIfAbsent(descriptor.runId(), ignored -> new CopyOnWriteArrayList<>())
|
||||
.add(subscriber);
|
||||
Runnable remove = () -> remove(descriptor.runId(), subscriber);
|
||||
emitter.onCompletion(remove);
|
||||
emitter.onTimeout(remove);
|
||||
emitter.onError(error -> remove.run());
|
||||
return emitter;
|
||||
}
|
||||
|
||||
/**
|
||||
* 周期性向全部本地订阅者重放新增事件并发送保活注释。
|
||||
*/
|
||||
@Scheduled(fixedDelayString = "${easyflow.agent.runtime.agui-replay-poll-interval:100ms}")
|
||||
public void dispatch() {
|
||||
for (Map.Entry<String, CopyOnWriteArrayList<Subscriber>> entry : subscribers.entrySet()) {
|
||||
dispatch(entry.getKey(), entry.getValue());
|
||||
}
|
||||
}
|
||||
|
||||
private void dispatch(String runId, CopyOnWriteArrayList<Subscriber> runSubscribers) {
|
||||
try {
|
||||
AgentAguiRunDescriptor descriptor = store.find(runId);
|
||||
if (descriptor == null) {
|
||||
runSubscribers.forEach(subscriber -> removeAndComplete(runId, subscriber));
|
||||
return;
|
||||
}
|
||||
Map<Long, List<Subscriber>> cursorGroups = new LinkedHashMap<>();
|
||||
for (Subscriber subscriber : runSubscribers) {
|
||||
cursorGroups.computeIfAbsent(subscriber.cursor, ignored -> new java.util.ArrayList<>())
|
||||
.add(subscriber);
|
||||
}
|
||||
for (Map.Entry<Long, List<Subscriber>> cursorGroup : cursorGroups.entrySet()) {
|
||||
long cursor = cursorGroup.getKey();
|
||||
List<String> events = store.readAfter(runId, cursor);
|
||||
for (Subscriber subscriber : cursorGroup.getValue()) {
|
||||
dispatchSubscriber(runId, subscriber, cursor, events, descriptor);
|
||||
}
|
||||
}
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("AG-UI 重连订阅调度失败: runId={}, subscriberCount={}",
|
||||
runId, runSubscribers.size(), exception);
|
||||
}
|
||||
}
|
||||
|
||||
private void dispatchSubscriber(String runId,
|
||||
Subscriber subscriber,
|
||||
long afterCursor,
|
||||
List<String> events,
|
||||
AgentAguiRunDescriptor descriptor) {
|
||||
try {
|
||||
for (int index = 0; index < events.size(); index++) {
|
||||
long cursor = afterCursor + index + 1L;
|
||||
subscriber.emitter.send(SseEmitter.event().id(Long.toString(cursor)).data(events.get(index)));
|
||||
subscriber.cursor = cursor;
|
||||
subscriber.lastWriteAt = System.currentTimeMillis();
|
||||
}
|
||||
if (descriptor.status().isTerminal() && subscriber.cursor >= descriptor.lastCursor()) {
|
||||
removeAndComplete(runId, subscriber);
|
||||
return;
|
||||
}
|
||||
long now = System.currentTimeMillis();
|
||||
if (events.isEmpty() && now - subscriber.lastWriteAt >= KEEP_ALIVE_INTERVAL_MILLIS) {
|
||||
subscriber.emitter.send(SseEmitter.event().comment("keepalive"));
|
||||
subscriber.lastWriteAt = now;
|
||||
}
|
||||
} catch (IOException | IllegalStateException exception) {
|
||||
LOG.debug("移除已断开的 AG-UI 重连订阅: runId={}, cursor={}",
|
||||
runId, subscriber.cursor, exception);
|
||||
remove(runId, subscriber);
|
||||
}
|
||||
}
|
||||
|
||||
private void removeAndComplete(String runId, Subscriber subscriber) {
|
||||
remove(runId, subscriber);
|
||||
subscriber.emitter.complete();
|
||||
}
|
||||
|
||||
private void remove(String runId, Subscriber subscriber) {
|
||||
CopyOnWriteArrayList<Subscriber> values = subscribers.get(runId);
|
||||
if (values == null) {
|
||||
return;
|
||||
}
|
||||
values.remove(subscriber);
|
||||
if (values.isEmpty()) {
|
||||
subscribers.remove(runId, values);
|
||||
}
|
||||
}
|
||||
|
||||
private static final class Subscriber {
|
||||
private final SseEmitter emitter;
|
||||
private long cursor;
|
||||
private long lastWriteAt = System.currentTimeMillis();
|
||||
|
||||
private Subscriber(SseEmitter emitter, long cursor) {
|
||||
this.emitter = emitter;
|
||||
this.cursor = cursor;
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -4,7 +4,7 @@ package tech.easyflow.agent.runtime.agui;
|
||||
* 单次 AG-UI 连接的客户端 wire 标识。
|
||||
*
|
||||
* @param threadId 客户端 thread ID
|
||||
* @param runId 客户端 run ID，仅用于协议输出
|
||||
* @param runId 客户端 run ID，用于协议输出、幂等查找与刷新重连
|
||||
* @param userMessageId 客户端本轮用户消息 ID
|
||||
* @param userMessageContent 客户端本轮用户消息正文
|
||||
*/
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
import tech.easyflow.core.chat.protocol.sse.ChatSseEmitter;
|
||||
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
/**
|
||||
* 将 AG-UI 运行输出写入可重放日志，同时尽力推送到首次 HTTP 连接。
|
||||
*
|
||||
* <p>底层浏览器连接关闭后，本对象仍保持逻辑可写，模型运行只会因终态或显式取消而结束。</p>
|
||||
*/
|
||||
public final class ResumableAguiSseEmitter extends ChatSseEmitter {
|
||||
|
||||
private static final Logger LOG = LoggerFactory.getLogger(ResumableAguiSseEmitter.class);
|
||||
|
||||
private final String threadId;
|
||||
private final String runId;
|
||||
private final AgentAguiRunJournal journal;
|
||||
private final SseEmitter subscriberEmitter;
|
||||
private final AtomicBoolean completed = new AtomicBoolean(false);
|
||||
private final AtomicBoolean journalFailed = new AtomicBoolean(false);
|
||||
|
||||
/**
|
||||
* 创建可重连 AG-UI SSE 发射器。
|
||||
*
|
||||
* @param threadId AG-UI 线程 ID
|
||||
* @param runId 运行 ID
|
||||
* @param journal 可重放事件日志
|
||||
* @param subscriberEmitter 从游标零开始的首次订阅
|
||||
*/
|
||||
public ResumableAguiSseEmitter(String threadId,
|
||||
String runId,
|
||||
AgentAguiRunJournal journal,
|
||||
SseEmitter subscriberEmitter) {
|
||||
super(0L);
|
||||
this.threadId = threadId;
|
||||
this.runId = runId;
|
||||
this.journal = journal;
|
||||
this.subscriberEmitter = java.util.Objects.requireNonNull(
|
||||
subscriberEmitter, "subscriberEmitter cannot be null");
|
||||
}
|
||||
|
||||
/**
|
||||
* 返回与重连相同的 Redis 游标订阅，首次连接不走旁路直推。
|
||||
*
|
||||
* @return 首次订阅 SSE
|
||||
*/
|
||||
@Override
|
||||
public SseEmitter getEmitter() {
|
||||
return subscriberEmitter;
|
||||
}
|
||||
|
||||
/**
|
||||
* 先记录事件，再尽力发送给当前浏览器连接。
|
||||
*
|
||||
* @param data 已编码 AG-UI 事件
|
||||
* @return 日志已接受事件时为 true
|
||||
*/
|
||||
@Override
|
||||
public boolean sendData(String data) {
|
||||
if (completed.get()) {
|
||||
return false;
|
||||
}
|
||||
try {
|
||||
journal.append(runId, data);
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("AG-UI 事件日志拒绝写入: runId={}", runId, exception);
|
||||
journalFailed.set(true);
|
||||
completed.set(true);
|
||||
failJournal(exception);
|
||||
super.complete();
|
||||
return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* 刷入尾部事件并关闭首次 HTTP 连接。
|
||||
*/
|
||||
@Override
|
||||
public void complete() {
|
||||
if (!completed.compareAndSet(false, true)) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
journal.complete(runId);
|
||||
} catch (RuntimeException exception) {
|
||||
LOG.error("AG-UI 运行终态日志同步写入失败，改写为日志故障终态: runId={}", runId, exception);
|
||||
journalFailed.set(true);
|
||||
failJournal(exception);
|
||||
} finally {
|
||||
super.complete();
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 刷入尾部事件并关闭首次 HTTP 连接。
|
||||
*
|
||||
* @param error 关闭原因
|
||||
*/
|
||||
@Override
|
||||
public void completeWithError(Throwable error) {
|
||||
complete();
|
||||
}
|
||||
|
||||
/**
|
||||
* 返回运行输出逻辑关闭状态，不受首次浏览器连接断开影响。
|
||||
*
|
||||
* @return 运行已经收口时为 true
|
||||
*/
|
||||
@Override
|
||||
public boolean isClosed() {
|
||||
return completed.get();
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断本次运行是否以可持久恢复的方式完成。
|
||||
*
|
||||
* @return 事件日志没有发生写入故障时为 {@code true}
|
||||
*/
|
||||
public boolean isJournalCompletionSuccessful() {
|
||||
return completed.get() && !journalFailed.get();
|
||||
}
|
||||
|
||||
private void failJournal(RuntimeException exception) {
|
||||
String detail = exception.getMessage() == null ? "" : exception.getMessage();
|
||||
boolean capacityFailure = detail.contains("安全上限") || detail.contains("16 MiB");
|
||||
String code = capacityFailure
|
||||
? "AGENT_RUN_JOURNAL_LIMIT"
|
||||
: "AGENT_RUN_JOURNAL_UNAVAILABLE";
|
||||
String message = capacityFailure
|
||||
? "Agent 输出超过可恢复日志容量，运行已停止"
|
||||
: "Agent 输出日志暂时不可用，运行已停止";
|
||||
try {
|
||||
journal.fail(threadId, runId, code, message);
|
||||
} catch (RuntimeException terminalException) {
|
||||
LOG.error("AG-UI 日志故障终态暂未写入，将由后台继续重试: runId={}",
|
||||
runId, terminalException);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -46,6 +46,15 @@ public interface AgentRunOutput {
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* 判断浏览器 SSE 连接断开时是否应取消底层 Agent 运行。
|
||||
*
|
||||
* @return 连接断开需要取消运行时为 true
|
||||
*/
|
||||
default boolean cancelRunOnDisconnect() {
|
||||
return true;
|
||||
}
|
||||
|
||||
/**
|
||||
* 发送协议终态并关闭连接。
|
||||
*
|
||||
|
||||
@@ -8,6 +8,7 @@ import com.easyagents.agui.AguiRuntimeEventProjector;
|
||||
import io.agentscope.core.agui.event.AguiEvent;
|
||||
import io.agentscope.core.agui.model.AguiMessage;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
import tech.easyflow.agent.runtime.agui.ResumableAguiSseEmitter;
|
||||
import tech.easyflow.agent.runtime.hitl.ToolApprovalInputProjection;
|
||||
import tech.easyflow.core.chat.protocol.ChatDomain;
|
||||
import tech.easyflow.core.chat.protocol.ChatType;
|
||||
@@ -35,6 +36,7 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
private final String clientUserMessageId;
|
||||
private final String clientUserMessageContent;
|
||||
private final ChatSseEmitter delegate;
|
||||
private final boolean cancelRunOnDisconnect;
|
||||
private final AguiRuntimeEventProjector projector;
|
||||
private final AguiProtocolEventEncoder encoder = new AguiProtocolEventEncoder();
|
||||
|
||||
@@ -56,7 +58,7 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
* @param clientUserMessageId 本轮客户端用户消息 ID
|
||||
*/
|
||||
public AguiAgentRunOutput(String threadId, String runId, String clientUserMessageId) {
|
||||
this(threadId, runId, clientUserMessageId, null, new ChatSseEmitter());
|
||||
this(threadId, runId, clientUserMessageId, null, new ChatSseEmitter(), true);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -72,7 +74,7 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
String runId,
|
||||
String clientUserMessageId,
|
||||
String clientUserMessageContent) {
|
||||
this(threadId, runId, clientUserMessageId, clientUserMessageContent, new ChatSseEmitter());
|
||||
this(threadId, runId, clientUserMessageId, clientUserMessageContent, new ChatSseEmitter(), true);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -88,7 +90,7 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
String runId,
|
||||
String clientUserMessageId,
|
||||
ChatSseEmitter delegate) {
|
||||
this(threadId, runId, clientUserMessageId, null, delegate);
|
||||
this(threadId, runId, clientUserMessageId, null, delegate, true);
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -106,11 +108,32 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
String clientUserMessageId,
|
||||
String clientUserMessageContent,
|
||||
ChatSseEmitter delegate) {
|
||||
this(threadId, runId, clientUserMessageId, clientUserMessageContent, delegate, true);
|
||||
}
|
||||
|
||||
/**
|
||||
* 使用指定连接生命周期策略创建 AG-UI 输出。
|
||||
*
|
||||
* @param threadId 客户端 thread ID
|
||||
* @param runId 客户端 run ID
|
||||
* @param clientUserMessageId 本轮客户端用户消息 ID
|
||||
* @param clientUserMessageContent 本轮客户端用户消息正文
|
||||
* @param delegate SSE 发射器
|
||||
* @param cancelRunOnDisconnect 浏览器连接断开时是否取消运行
|
||||
*/
|
||||
public AguiAgentRunOutput(
|
||||
String threadId,
|
||||
String runId,
|
||||
String clientUserMessageId,
|
||||
String clientUserMessageContent,
|
||||
ChatSseEmitter delegate,
|
||||
boolean cancelRunOnDisconnect) {
|
||||
this.threadId = requireText(threadId, "threadId");
|
||||
this.runId = requireText(runId, "runId");
|
||||
this.clientUserMessageId = clientUserMessageId;
|
||||
this.clientUserMessageContent = clientUserMessageContent;
|
||||
this.delegate = java.util.Objects.requireNonNull(delegate, "delegate cannot be null");
|
||||
this.cancelRunOnDisconnect = cancelRunOnDisconnect;
|
||||
this.projector = new AguiRuntimeEventProjector(threadId, runId);
|
||||
}
|
||||
|
||||
@@ -119,6 +142,15 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
return delegate.getEmitter();
|
||||
}
|
||||
|
||||
/**
|
||||
* 获取公开 AG-UI 运行 ID。
|
||||
*
|
||||
* @return 运行 ID
|
||||
*/
|
||||
public String runId() {
|
||||
return runId;
|
||||
}
|
||||
|
||||
@Override
|
||||
public synchronized boolean emitRuntimeEvent(AgentRuntimeEvent event) {
|
||||
if (event == null || event.getEventType() == null || delegate.isClosed()) {
|
||||
@@ -249,6 +281,11 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
return pendingCompletedEvent != null;
|
||||
}
|
||||
|
||||
@Override
|
||||
public boolean cancelRunOnDisconnect() {
|
||||
return cancelRunOnDisconnect;
|
||||
}
|
||||
|
||||
@Override
|
||||
public synchronized boolean finish(String finalText) {
|
||||
if (delegate.isClosed()) {
|
||||
@@ -277,6 +314,9 @@ public final class AguiAgentRunOutput implements AgentRunOutput {
|
||||
}
|
||||
}
|
||||
delegate.complete();
|
||||
if (delegate instanceof ResumableAguiSseEmitter resumableEmitter) {
|
||||
return resumableEmitter.isJournalCompletionSuccessful();
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
|
||||
@@ -164,6 +164,31 @@ public class AgentRuntimeCommandConsumerTest {
|
||||
Assert.assertEquals("cmd-cancel", resultRegistry.lastSuccessCommandId);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证跨节点精确取消命令保留请求、用户和原因。
|
||||
*
|
||||
* @throws Exception 消息序列化异常
|
||||
*/
|
||||
@Test
|
||||
public void consumerShouldHandleCancelRunCommand() throws Exception {
|
||||
AgentRuntimeProperties properties = new AgentRuntimeProperties();
|
||||
properties.setInstanceId("node-a");
|
||||
RecordingAgentRunService service = new RecordingAgentRunService();
|
||||
RecordingCommandResultRegistry resultRegistry = new RecordingCommandResultRegistry();
|
||||
AgentRuntimeCommandConsumer consumer = new AgentRuntimeCommandConsumer(
|
||||
new ObjectMapper(), properties, new MQProperties(), service, resultRegistry);
|
||||
AgentRuntimeCommandMessage command = command("cmd-cancel-run", "node-a");
|
||||
command.setAction(AgentRuntimeCommandAction.CANCEL_RUN);
|
||||
command.setReason("用户停止");
|
||||
|
||||
consumer.handle(List.of(message(command)));
|
||||
|
||||
Assert.assertEquals("request-cmd-cancel-run", service.lastCancelledRequestId);
|
||||
Assert.assertEquals("1", service.lastCancelledUserId);
|
||||
Assert.assertEquals("用户停止", service.lastReason);
|
||||
Assert.assertEquals("cmd-cancel-run", resultRegistry.lastSuccessCommandId);
|
||||
}
|
||||
|
||||
private AgentRuntimeCommandMessage command(String commandId, String targetNodeId) {
|
||||
AgentRuntimeCommandMessage command = new AgentRuntimeCommandMessage();
|
||||
command.setCommandId(commandId);
|
||||
@@ -191,6 +216,8 @@ public class AgentRuntimeCommandConsumerTest {
|
||||
private String lastRequestId;
|
||||
private String lastReason;
|
||||
private String lastCancelledAgentId;
|
||||
private String lastCancelledRequestId;
|
||||
private String lastCancelledUserId;
|
||||
private String lastApprovalId;
|
||||
|
||||
@Override
|
||||
@@ -221,6 +248,13 @@ public class AgentRuntimeCommandConsumerTest {
|
||||
public void cancelAgentLocal(String agentId) {
|
||||
lastCancelledAgentId = agentId;
|
||||
}
|
||||
|
||||
@Override
|
||||
public void cancelRunLocal(String requestId, String userId, String reason) {
|
||||
lastCancelledRequestId = requestId;
|
||||
lastCancelledUserId = userId;
|
||||
lastReason = reason;
|
||||
}
|
||||
}
|
||||
|
||||
private static class RecordingCommandResultRegistry extends AgentRuntimeCommandResultRegistry {
|
||||
|
||||
@@ -4,9 +4,13 @@ import com.easyagents.agent.runtime.AgentInitRequest;
|
||||
import com.easyagents.agent.runtime.AgentResumeRequest;
|
||||
import com.easyagents.agent.runtime.AgentRuntime;
|
||||
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
|
||||
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import reactor.core.Disposable;
|
||||
import reactor.core.publisher.Flux;
|
||||
import tech.easyflow.agent.distributed.AgentRuntimeRouteRegistry;
|
||||
import tech.easyflow.agent.runtime.lock.AgentRunLock;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
import tech.easyflow.core.runtime.ChatAssistantAccumulator;
|
||||
import tech.easyflow.core.runtime.ChatRuntimeContext;
|
||||
@@ -14,6 +18,10 @@ import tech.easyflow.core.runtime.ChatRuntimeContext;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
import java.util.concurrent.atomic.AtomicReference;
|
||||
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.doThrow;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
/**
|
||||
* Agent 运行态注册表测试。
|
||||
*/
|
||||
@@ -69,6 +77,39 @@ public class AgentRunRegistryTest {
|
||||
Assert.assertThrows(BusinessException.class, () -> registry.approve(null, "token-3", "user-1"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证分布式锁释放失败时仍会关闭 Runtime 并清理全部本地、远程索引。
|
||||
*
|
||||
* @throws Exception Runtime close 方法声明的受检异常
|
||||
*/
|
||||
@Test
|
||||
public void removeShouldFinishCleanupWhenLockReleaseFails() throws Exception {
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
AgentRuntimeRouteRegistry routeRegistry = mock(AgentRuntimeRouteRegistry.class);
|
||||
AgentRuntime runtime = mock(AgentRuntime.class);
|
||||
AgentRunLock.Handle lockHandle = mock(AgentRunLock.Handle.class);
|
||||
doThrow(new IllegalStateException("redis unavailable")).when(lockHandle).release();
|
||||
registry.setRouteRegistry(routeRegistry);
|
||||
AgentRunRegistry.AgentRunContext context = new AgentRunRegistry.AgentRunContext(
|
||||
"request-cleanup", "session-cleanup", runtime, null,
|
||||
new ChatRuntimeContext(), new StringBuilder(), new ChatAssistantAccumulator(),
|
||||
new AtomicBoolean(false), false,
|
||||
new AgentRunRegistry.RunOwner("agent-1", "session-cleanup", "user-1"),
|
||||
lockHandle, event -> { }, error -> { }, () -> { });
|
||||
registry.register(context);
|
||||
registry.registerResumeToken("request-cleanup", "token-cleanup");
|
||||
String approvalId = registry.registerApproval("request-cleanup", "token-cleanup");
|
||||
|
||||
Assert.assertThrows(IllegalStateException.class, () -> registry.remove("request-cleanup"));
|
||||
|
||||
verify(runtime).close();
|
||||
verify(routeRegistry).removeResumeToken("token-cleanup");
|
||||
verify(routeRegistry).removeApproval(approvalId);
|
||||
verify(routeRegistry).removeRun("request-cleanup");
|
||||
Assert.assertNull(registry.get("request-cleanup"));
|
||||
Assert.assertThrows(BusinessException.class, () -> registry.resolveApproval(approvalId));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证运行审批只能由运行发起人处理。
|
||||
*/
|
||||
@@ -144,6 +185,85 @@ public class AgentRunRegistryTest {
|
||||
Assert.assertNotNull(registry.get("request-10"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证显式单次取消会校验归属并进入标准取消事件链。
|
||||
*/
|
||||
@Test
|
||||
public void cancelRunShouldEmitNormalizedCancellationForOwner() {
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
AtomicReference<AgentRuntimeEvent> cancellation = new AtomicReference<>();
|
||||
AgentRunRegistry.AgentRunContext context = new AgentRunRegistry.AgentRunContext(
|
||||
"request-11",
|
||||
"session-11",
|
||||
new CapturingRuntime(),
|
||||
null,
|
||||
new ChatRuntimeContext(),
|
||||
new StringBuilder(),
|
||||
new ChatAssistantAccumulator(),
|
||||
new AtomicBoolean(false),
|
||||
false,
|
||||
new AgentRunRegistry.RunOwner("agent-1", "session-11", "user-1"),
|
||||
null,
|
||||
cancellation::set,
|
||||
error -> {
|
||||
},
|
||||
() -> {
|
||||
});
|
||||
registry.register(context);
|
||||
|
||||
Assert.assertThrows(BusinessException.class,
|
||||
() -> registry.cancelRun("request-11", "user-2", "用户停止"));
|
||||
registry.cancelRun("request-11", "user-1", "用户停止");
|
||||
|
||||
Assert.assertNotNull(cancellation.get());
|
||||
Assert.assertEquals(AgentRuntimeEventType.CANCELLED, cancellation.get().getEventType());
|
||||
Assert.assertEquals("用户停止", cancellation.get().getPayload().get("reason"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证自然终态或重复远程命令到达后，显式取消保持幂等。
|
||||
*/
|
||||
@Test
|
||||
public void cancelRunShouldBeIdempotentAfterContextRemoved() {
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
|
||||
registry.cancelRun("request-finished", "user-1", "用户停止");
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证订阅启动后上下文已经移除时，迟到的 Disposable 会立即释放。
|
||||
*/
|
||||
@Test
|
||||
public void bindSubscriptionShouldDisposeWhenContextAlreadyRemoved() {
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
Disposable subscription = mock(Disposable.class);
|
||||
|
||||
registry.bindSubscription("request-removed", subscription);
|
||||
|
||||
verify(subscription).dispose();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证运行已收口时绑定订阅不会遗留后台模型流。
|
||||
*/
|
||||
@Test
|
||||
public void bindSubscriptionShouldDisposeWhenContextAlreadyFinished() {
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
AtomicBoolean finished = new AtomicBoolean(true);
|
||||
AgentRunRegistry.AgentRunContext context = new AgentRunRegistry.AgentRunContext(
|
||||
"request-finished", "session-finished", new CapturingRuntime(), null,
|
||||
new ChatRuntimeContext(), new StringBuilder(), new ChatAssistantAccumulator(),
|
||||
finished, false,
|
||||
new AgentRunRegistry.RunOwner("agent-1", "session-finished", "user-1"),
|
||||
null, event -> { }, error -> { }, () -> { });
|
||||
Disposable subscription = mock(Disposable.class);
|
||||
registry.register(context);
|
||||
|
||||
registry.bindSubscription("request-finished", subscription);
|
||||
|
||||
verify(subscription).dispose();
|
||||
}
|
||||
|
||||
private AgentRunRegistry.AgentRunContext context(String requestId,
|
||||
String sessionId,
|
||||
String userId,
|
||||
|
||||
@@ -16,6 +16,7 @@ import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.Mockito;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
import tech.easyflow.agent.entity.AgentHitlPending;
|
||||
import tech.easyflow.agent.entity.Agent;
|
||||
import tech.easyflow.agent.entity.AgentKnowledgeBinding;
|
||||
@@ -24,6 +25,12 @@ import tech.easyflow.agent.distributed.AgentRuntimeCommandProducer;
|
||||
import tech.easyflow.agent.distributed.AgentRuntimeRoute;
|
||||
import tech.easyflow.agent.distributed.AgentRuntimeRouteRegistry;
|
||||
import tech.easyflow.agent.runtime.event.AgentRunEventRecorder;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunDescriptor;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunJournal;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunStatus;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunStore;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunSubscriptionService;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiWireContext;
|
||||
import tech.easyflow.agent.runtime.hitl.AgentHitlPendingService;
|
||||
import tech.easyflow.agent.runtime.document.AgentDocumentContext;
|
||||
import tech.easyflow.agent.runtime.lock.AgentRunLock;
|
||||
@@ -50,6 +57,11 @@ import java.lang.reflect.Method;
|
||||
import java.math.BigInteger;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicBoolean;
|
||||
|
||||
/**
|
||||
@@ -57,6 +69,154 @@ import java.util.concurrent.atomic.AtomicBoolean;
|
||||
*/
|
||||
public class AgentRunServiceDraftAndHitlTest {
|
||||
|
||||
/**
|
||||
* 创建用于 owner 恢复测试的运行描述。
|
||||
*
|
||||
* @param runId 运行 ID
|
||||
* @param requestId 请求 ID
|
||||
* @param updatedAt 最近更新时间
|
||||
* @return 运行描述
|
||||
*/
|
||||
private static AgentAguiRunDescriptor descriptor(String runId, String requestId, long updatedAt) {
|
||||
return new AgentAguiRunDescriptor(
|
||||
runId, requestId, "thread-" + runId, "7", "session-" + runId,
|
||||
"1", "2", false, AgentAguiRunStatus.RUNNING, 0L, updatedAt, updatedAt);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证可重连运行会在 Redis store 创建前注册临时 owner 路由。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void createRunOutputShouldRegisterProvisionalOwnerBeforeStore() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
AgentRuntimeRouteRegistry routeRegistry = Mockito.mock(AgentRuntimeRouteRegistry.class);
|
||||
AgentAguiRunStore store = Mockito.mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = Mockito.mock(AgentAguiRunJournal.class);
|
||||
AgentAguiRunSubscriptionService subscriptions = Mockito.mock(AgentAguiRunSubscriptionService.class);
|
||||
Mockito.when(subscriptions.subscribe(Mockito.any(AgentAguiRunDescriptor.class), Mockito.eq(0L)))
|
||||
.thenReturn(new SseEmitter());
|
||||
setField(service, "agentRuntimeRouteRegistry", routeRegistry);
|
||||
setField(service, "agentAguiRunStore", store);
|
||||
setField(service, "agentAguiRunJournal", journal);
|
||||
setField(service, "agentAguiRunSubscriptionService", subscriptions);
|
||||
Agent agent = new Agent();
|
||||
agent.setId(BigInteger.valueOf(7L));
|
||||
LoginAccount account = new LoginAccount();
|
||||
account.setId(BigInteger.ONE);
|
||||
account.setTenantId(BigInteger.TWO);
|
||||
|
||||
AgentRunOutput output = invoke(service, "createRunOutput",
|
||||
new Class<?>[]{AgentAguiWireContext.class, String.class, LoginAccount.class,
|
||||
Agent.class, String.class, boolean.class},
|
||||
new AgentAguiWireContext("thread-start", "run-start", "user-start", "启动"),
|
||||
"request-start", account, agent, "session-start", false);
|
||||
|
||||
Assert.assertNotNull(output);
|
||||
org.mockito.InOrder order = Mockito.inOrder(routeRegistry, store, subscriptions);
|
||||
order.verify(routeRegistry).registerRun("request-start", "7");
|
||||
order.verify(store).create(Mockito.any(AgentAguiRunDescriptor.class));
|
||||
order.verify(subscriptions).subscribe(Mockito.any(AgentAguiRunDescriptor.class), Mockito.eq(0L));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证会话锁获取失败时会写入失败终态并移除临时 owner 路由。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void runAguiAwareShouldCloseProvisionalRunWhenLockAcquireFails() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
AgentRunLock runLock = Mockito.mock(AgentRunLock.class);
|
||||
AgentRuntimeRouteRegistry routeRegistry = Mockito.mock(AgentRuntimeRouteRegistry.class);
|
||||
AgentRunOutput output = Mockito.mock(AgentRunOutput.class);
|
||||
Mockito.when(runLock.acquire(Mockito.any(), Mockito.eq("session-lock-failed")))
|
||||
.thenThrow(new IllegalStateException("lock unavailable"));
|
||||
Mockito.when(output.cancelRunOnDisconnect()).thenReturn(false);
|
||||
setField(service, "agentRunLock", runLock);
|
||||
setField(service, "agentRuntimeRouteRegistry", routeRegistry);
|
||||
Agent agent = new Agent();
|
||||
agent.setId(BigInteger.valueOf(8L));
|
||||
|
||||
try {
|
||||
invoke(service, "runAguiAware",
|
||||
new Class<?>[]{Agent.class, String.class, List.class, List.class,
|
||||
LoginAccount.class, String.class, String.class, String.class,
|
||||
String.class, ChatRuntimeContext.class, boolean.class,
|
||||
AgentSessionStore.class, AgentRunOutput.class},
|
||||
agent, "问题", List.of(), List.of(), new LoginAccount(),
|
||||
"request-lock-failed", "trace-lock-failed", "session-lock-failed",
|
||||
"AGENT", new ChatRuntimeContext(), true,
|
||||
Mockito.mock(AgentSessionStore.class), output);
|
||||
Assert.fail("expected lock failure");
|
||||
} catch (Exception exception) {
|
||||
Assert.assertTrue(rootCause(exception) instanceof IllegalStateException);
|
||||
}
|
||||
|
||||
Mockito.verify(output).completeWithError(Mockito.any(IllegalStateException.class));
|
||||
Mockito.verify(routeRegistry).removeRun("request-lock-failed");
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证超过启动宽限期但 owner 节点仍存活时不会误判运行丢失。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void reconcileShouldKeepSlowStartingRunWhenOwnerBootIsAlive() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
AgentRuntimeRouteRegistry routeRegistry = Mockito.mock(AgentRuntimeRouteRegistry.class);
|
||||
AgentAguiRunStore store = Mockito.mock(AgentAguiRunStore.class);
|
||||
AgentRuntimeRoute route = new AgentRuntimeRoute();
|
||||
route.setNodeId("node-alive");
|
||||
route.setBootId("boot-alive");
|
||||
Mockito.when(routeRegistry.findOwnerRoute("request-slow")).thenReturn(route);
|
||||
Mockito.when(routeRegistry.currentNodeBootId("node-alive")).thenReturn("boot-alive");
|
||||
setField(service, "agentRuntimeRouteRegistry", routeRegistry);
|
||||
setField(service, "agentAguiRunStore", store);
|
||||
long old = System.currentTimeMillis() - 120_000L;
|
||||
AgentAguiRunDescriptor descriptor = descriptor("run-slow", "request-slow", old);
|
||||
|
||||
AgentAguiRunDescriptor result = invoke(service, "reconcileInterruptedRun",
|
||||
new Class<?>[]{AgentAguiRunDescriptor.class}, descriptor);
|
||||
|
||||
Assert.assertSame(descriptor, result);
|
||||
Mockito.verify(store, Mockito.never()).failOwnerLost(Mockito.any());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 owner 启动代心跳消失且超过宽限期后写入明确失败终态。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void reconcileShouldFailRunWhenOwnerBootDisappears() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
AgentRuntimeRouteRegistry routeRegistry = Mockito.mock(AgentRuntimeRouteRegistry.class);
|
||||
AgentAguiRunStore store = Mockito.mock(AgentAguiRunStore.class);
|
||||
AgentRuntimeRoute route = new AgentRuntimeRoute();
|
||||
route.setNodeId("node-lost");
|
||||
route.setBootId("boot-lost");
|
||||
Mockito.when(routeRegistry.findOwnerRoute("request-lost")).thenReturn(route);
|
||||
Mockito.when(routeRegistry.currentNodeBootId("node-lost")).thenReturn(null);
|
||||
long old = System.currentTimeMillis() - 120_000L;
|
||||
AgentAguiRunDescriptor descriptor = descriptor("run-lost", "request-lost", old);
|
||||
AgentAguiRunDescriptor failed = new AgentAguiRunDescriptor(
|
||||
descriptor.runId(), descriptor.requestId(), descriptor.threadId(), descriptor.agentId(),
|
||||
descriptor.sessionId(), descriptor.userId(), descriptor.tenantId(), descriptor.draft(),
|
||||
AgentAguiRunStatus.FAILED, 1L, descriptor.createdAt(), System.currentTimeMillis());
|
||||
Mockito.when(store.find("run-lost")).thenReturn(failed);
|
||||
setField(service, "agentRuntimeRouteRegistry", routeRegistry);
|
||||
setField(service, "agentAguiRunStore", store);
|
||||
|
||||
AgentAguiRunDescriptor result = invoke(service, "reconcileInterruptedRun",
|
||||
new Class<?>[]{AgentAguiRunDescriptor.class}, descriptor);
|
||||
|
||||
Assert.assertSame(failed, result);
|
||||
Mockito.verify(store).failOwnerLost(descriptor);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证工具 HITL 事件会映射为显式前端载荷。
|
||||
*
|
||||
@@ -418,6 +578,78 @@ public class AgentRunServiceDraftAndHitlTest {
|
||||
Mockito.verify(output).complete();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证可恢复终态写入失败时不会把聊天会话记录为完成。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void finishClaimedRunShouldRecordFailureWhenOutputCompletionFails() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
RecordingChatRuntimeManager chatRuntimeManager = new RecordingChatRuntimeManager();
|
||||
setField(service, "agentRunRegistry", new AgentRunRegistry());
|
||||
setField(service, "chatRuntimeManager", chatRuntimeManager);
|
||||
AgentRunOutput output = Mockito.mock(AgentRunOutput.class);
|
||||
Mockito.when(output.finish("最终正文")).thenReturn(false);
|
||||
|
||||
invoke(service, "finishClaimedRun",
|
||||
new Class<?>[]{String.class, AgentRunOutput.class, ChatRuntimeContext.class,
|
||||
StringBuilder.class, ChatAssistantAccumulator.class, boolean.class, List.class},
|
||||
"request-output-failed", output, chatContext(), new StringBuilder("最终正文"),
|
||||
new ChatAssistantAccumulator(), true, List.of());
|
||||
|
||||
Assert.assertEquals(1, chatRuntimeManager.recordAssistantCompletedCount);
|
||||
Assert.assertEquals(1, chatRuntimeManager.recordFailureCount);
|
||||
Assert.assertEquals(0, chatRuntimeManager.recordCompletedCount);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证终态回调内部抛错后强制释放 Registry 并写入失败状态。
|
||||
*
|
||||
* @throws Exception 反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void runtimeCallbackFailureShouldForceCleanupAfterTerminalClaim() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
AgentRunRegistry registry = new AgentRunRegistry();
|
||||
RecordingChatRuntimeManager chatRuntimeManager = new RecordingChatRuntimeManager();
|
||||
AgentRunOutput output = Mockito.mock(AgentRunOutput.class);
|
||||
AtomicBoolean finished = new AtomicBoolean(true);
|
||||
AgentRunRegistry.AgentRunContext context = new AgentRunRegistry.AgentRunContext(
|
||||
"request-callback-failed",
|
||||
"session-callback-failed",
|
||||
new NoopRuntime(),
|
||||
output,
|
||||
chatContext(),
|
||||
new StringBuilder(),
|
||||
new ChatAssistantAccumulator(),
|
||||
finished,
|
||||
true,
|
||||
new AgentRunRegistry.RunOwner("agent-1", "session-callback-failed", "user-1"),
|
||||
null,
|
||||
event -> {
|
||||
},
|
||||
error -> {
|
||||
},
|
||||
() -> {
|
||||
}
|
||||
);
|
||||
registry.register(context);
|
||||
setField(service, "agentRunRegistry", registry);
|
||||
setField(service, "chatRuntimeManager", chatRuntimeManager);
|
||||
|
||||
invoke(service, "runRuntimeCallbackSafely",
|
||||
new Class<?>[]{Runnable.class, String.class, AgentRunOutput.class,
|
||||
ChatRuntimeContext.class, AtomicBoolean.class, boolean.class},
|
||||
(Runnable) () -> {
|
||||
throw new IllegalStateException("审计写入失败");
|
||||
}, "request-callback-failed", output, chatContext(), finished, true);
|
||||
|
||||
Assert.assertNull(registry.get("request-callback-failed"));
|
||||
Assert.assertEquals(1, chatRuntimeManager.recordFailureCount);
|
||||
Mockito.verify(output).completeWithError(Mockito.any(IllegalStateException.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证取消事件作为业务状态收口，不按系统错误发送。
|
||||
*
|
||||
@@ -456,6 +688,73 @@ public class AgentRunServiceDraftAndHitlTest {
|
||||
Assert.assertEquals(1, chatRuntimeManager.recordFailureCount);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证用户取消与自然完成并发到达时只允许一个终态完成协议投影和聊天记录收口。
|
||||
*
|
||||
* @throws Exception 并发任务或反射调用失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void concurrentCancelAndCompleteShouldShareOneTerminalArbitration() throws Exception {
|
||||
AgentRunService service = new AgentRunService();
|
||||
RecordingChatRuntimeManager chatRuntimeManager = new RecordingChatRuntimeManager();
|
||||
setField(service, "agentRunRegistry", new AgentRunRegistry());
|
||||
setField(service, "chatRuntimeManager", chatRuntimeManager);
|
||||
AgentRunOutput output = Mockito.mock(AgentRunOutput.class);
|
||||
Mockito.when(output.emitRuntimeEvent(Mockito.any())).thenReturn(true);
|
||||
Mockito.when(output.emitViewEvent(Mockito.any(), Mockito.any(), Mockito.any())).thenReturn(true);
|
||||
Mockito.when(output.finish(Mockito.anyString())).thenReturn(true);
|
||||
AtomicBoolean finished = new AtomicBoolean(false);
|
||||
CountDownLatch start = new CountDownLatch(1);
|
||||
ExecutorService executor = Executors.newFixedThreadPool(2);
|
||||
AgentRuntimeEvent completed = AgentRuntimeEvent.of(AgentRuntimeEventType.COMPLETED);
|
||||
completed.getPayload().put("text", "自然完成正文");
|
||||
AgentRuntimeEvent cancelled = AgentRuntimeEvent.of(AgentRuntimeEventType.CANCELLED);
|
||||
cancelled.getPayload().put("reason", "用户停止");
|
||||
|
||||
try {
|
||||
Future<?> completeFuture = executor.submit(() -> invokeTerminalEvent(
|
||||
service, start, completed, output, finished));
|
||||
Future<?> cancelFuture = executor.submit(() -> invokeTerminalEvent(
|
||||
service, start, cancelled, output, finished));
|
||||
start.countDown();
|
||||
completeFuture.get(5, TimeUnit.SECONDS);
|
||||
cancelFuture.get(5, TimeUnit.SECONDS);
|
||||
} finally {
|
||||
executor.shutdownNow();
|
||||
}
|
||||
|
||||
Assert.assertTrue(finished.get());
|
||||
Assert.assertEquals(1,
|
||||
chatRuntimeManager.recordCompletedCount + chatRuntimeManager.recordFailureCount);
|
||||
Mockito.verify(output, Mockito.times(1)).emitRuntimeEvent(Mockito.any());
|
||||
Mockito.verify(output, Mockito.times(1)).finish(Mockito.nullable(String.class));
|
||||
}
|
||||
|
||||
/**
|
||||
* 等待统一起跑信号后投递指定终态事件。
|
||||
*
|
||||
* @param service 被测服务
|
||||
* @param start 并发起跑信号
|
||||
* @param event 终态事件
|
||||
* @param output 运行输出
|
||||
* @param finished 共享终态标记
|
||||
*/
|
||||
private void invokeTerminalEvent(AgentRunService service,
|
||||
CountDownLatch start,
|
||||
AgentRuntimeEvent event,
|
||||
AgentRunOutput output,
|
||||
AtomicBoolean finished) {
|
||||
try {
|
||||
start.await(5, TimeUnit.SECONDS);
|
||||
invoke(service, "handleRuntimeEvent",
|
||||
runtimeEventParameterTypes(),
|
||||
event, "request-terminal-race", output, new StringBuilder("部分正文"),
|
||||
new ChatAssistantAccumulator(), chatContext(), finished, true);
|
||||
} catch (Exception e) {
|
||||
throw new IllegalStateException("并发终态测试执行失败", e);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证最终知识库引用会保留命中分片原文。
|
||||
*
|
||||
|
||||
@@ -0,0 +1,194 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
|
||||
import java.util.List;
|
||||
import java.util.concurrent.CountDownLatch;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.Future;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.concurrent.atomic.AtomicInteger;
|
||||
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.Mockito.doReturn;
|
||||
import static org.mockito.Mockito.doAnswer;
|
||||
import static org.mockito.Mockito.doThrow;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
/**
|
||||
* AG-UI 批量事件日志测试。
|
||||
*/
|
||||
public class AgentAguiRunJournalTest {
|
||||
|
||||
/**
|
||||
* 验证相邻同消息 delta 会在一次 Redis 批次中合并且顺序不变。
|
||||
*/
|
||||
@Test
|
||||
public void shouldMergeAdjacentDeltasBeforeRedisFlush() {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
|
||||
journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_CONTENT\",\"messageId\":\"m1\",\"delta\":\"你\"}");
|
||||
journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_CONTENT\",\"messageId\":\"m1\",\"delta\":\"好\"}");
|
||||
journal.append("run-1", "{\"type\":\"TEXT_MESSAGE_END\",\"messageId\":\"m1\"}");
|
||||
|
||||
journal.flush("run-1");
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<List<String>> events = ArgumentCaptor.forClass(List.class);
|
||||
verify(store).append(eq("run-1"), anyString(), events.capture());
|
||||
Assert.assertEquals(2, events.getValue().size());
|
||||
Assert.assertTrue(events.getValue().get(0).contains("\"delta\":\"你好\""));
|
||||
Assert.assertTrue(events.getValue().get(1).contains("TEXT_MESSAGE_END"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证同步终态写入失败后，后台刷新仍会保留并重试原批次。
|
||||
*/
|
||||
@Test
|
||||
public void shouldRetryTerminalBatchWithoutDroppingEvents() {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
|
||||
String terminal = "{\"type\":\"RUN_FINISHED\",\"runId\":\"run-1\",\"threadId\":\"1\"}";
|
||||
journal.append("run-1", terminal);
|
||||
doThrow(new IllegalStateException("redis down"))
|
||||
.doReturn(1L)
|
||||
.when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());
|
||||
|
||||
try {
|
||||
journal.complete("run-1");
|
||||
Assert.fail("首次同步写入应抛出异常");
|
||||
} catch (IllegalStateException expected) {
|
||||
Assert.assertEquals("redis down", expected.getMessage());
|
||||
}
|
||||
journal.flushPending();
|
||||
|
||||
ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
|
||||
verify(store, times(2)).append(
|
||||
eq("run-1"), batchIds.capture(), org.mockito.ArgumentMatchers.anyList());
|
||||
Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证结果不确定的批次内容保持冻结，期间到达的新事件使用新批次写入。
|
||||
*/
|
||||
@Test
|
||||
public void shouldKeepAmbiguousBatchImmutableAndFlushTailSeparately() {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
|
||||
doThrow(new IllegalStateException("ambiguous timeout"))
|
||||
.doReturn(1L)
|
||||
.doReturn(2L)
|
||||
.when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());
|
||||
String first = "{\"type\":\"RUN_STARTED\",\"runId\":\"run-1\"}";
|
||||
String tail = "{\"type\":\"TEXT_MESSAGE_START\",\"messageId\":\"m1\"}";
|
||||
journal.append("run-1", first);
|
||||
|
||||
Assert.assertThrows(IllegalStateException.class, () -> journal.flush("run-1"));
|
||||
journal.append("run-1", tail);
|
||||
journal.flushPending();
|
||||
|
||||
ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<List<String>> batches = ArgumentCaptor.forClass(List.class);
|
||||
verify(store, times(3)).append(eq("run-1"), batchIds.capture(), batches.capture());
|
||||
Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
|
||||
Assert.assertNotEquals(batchIds.getAllValues().get(1), batchIds.getAllValues().get(2));
|
||||
Assert.assertEquals(List.of(first), batches.getAllValues().get(0));
|
||||
Assert.assertEquals(List.of(first), batches.getAllValues().get(1));
|
||||
Assert.assertEquals(List.of(tail), batches.getAllValues().get(2));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证终态完成会等待在途后台刷盘，并在其失败后接管重试全部剩余批次。
|
||||
*
|
||||
* @throws Exception 并发任务等待失败时抛出
|
||||
*/
|
||||
@Test
|
||||
public void completeShouldWaitForBackgroundFlushAndPersistAllTailEvents() throws Exception {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
|
||||
CountDownLatch firstAppendStarted = new CountDownLatch(1);
|
||||
CountDownLatch releaseFirstAppend = new CountDownLatch(1);
|
||||
CountDownLatch completeStarted = new CountDownLatch(1);
|
||||
AtomicInteger appendCalls = new AtomicInteger();
|
||||
doAnswer(invocation -> {
|
||||
int call = appendCalls.incrementAndGet();
|
||||
if (call == 1) {
|
||||
firstAppendStarted.countDown();
|
||||
Assert.assertTrue(releaseFirstAppend.await(5, TimeUnit.SECONDS));
|
||||
throw new IllegalStateException("first append failed");
|
||||
}
|
||||
return (long) call;
|
||||
}).when(store).append(eq("run-1"), anyString(), org.mockito.ArgumentMatchers.anyList());
|
||||
String started = "{\"type\":\"RUN_STARTED\",\"runId\":\"run-1\"}";
|
||||
String finished = "{\"type\":\"RUN_FINISHED\",\"runId\":\"run-1\",\"threadId\":\"1\"}";
|
||||
journal.append("run-1", started);
|
||||
ExecutorService executor = Executors.newFixedThreadPool(2);
|
||||
|
||||
try {
|
||||
Future<?> background = executor.submit(journal::flushPending);
|
||||
Assert.assertTrue(firstAppendStarted.await(5, TimeUnit.SECONDS));
|
||||
journal.append("run-1", finished);
|
||||
Future<?> completion = executor.submit(() -> {
|
||||
completeStarted.countDown();
|
||||
journal.complete("run-1");
|
||||
});
|
||||
Assert.assertTrue(completeStarted.await(5, TimeUnit.SECONDS));
|
||||
Assert.assertFalse(completion.isDone());
|
||||
|
||||
releaseFirstAppend.countDown();
|
||||
background.get(5, TimeUnit.SECONDS);
|
||||
completion.get(5, TimeUnit.SECONDS);
|
||||
} finally {
|
||||
releaseFirstAppend.countDown();
|
||||
executor.shutdownNow();
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<List<String>> batches = ArgumentCaptor.forClass(List.class);
|
||||
ArgumentCaptor<String> batchIds = ArgumentCaptor.forClass(String.class);
|
||||
verify(store, times(3)).append(eq("run-1"), batchIds.capture(), batches.capture());
|
||||
Assert.assertEquals(batchIds.getAllValues().get(0), batchIds.getAllValues().get(1));
|
||||
Assert.assertEquals(List.of(started), batches.getAllValues().get(1));
|
||||
Assert.assertEquals(List.of(finished), batches.getAllValues().get(2));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证容量故障会丢弃未确认普通事件并以独立小批次写入失败终态。
|
||||
*/
|
||||
@Test
|
||||
public void shouldReplacePendingEventsWithJournalFailureTerminal() {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(store, new ObjectMapper());
|
||||
journal.append("run-1", "{\"type\":\"RUN_STARTED\"}");
|
||||
|
||||
journal.fail("thread-1", "run-1", "AGENT_RUN_JOURNAL_LIMIT", "输出过长");
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<List<String>> events = ArgumentCaptor.forClass(List.class);
|
||||
verify(store).append(eq("run-1"), anyString(), events.capture());
|
||||
Assert.assertEquals(1, events.getValue().size());
|
||||
Assert.assertTrue(events.getValue().get(0).contains("AGENT_RUN_JOURNAL_LIMIT"));
|
||||
Assert.assertFalse(events.getValue().get(0).contains("RUN_STARTED"));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证异常大的单事件在进入共享内存缓冲前即被拒绝。
|
||||
*/
|
||||
@Test
|
||||
public void shouldRejectOversizedSingleEvent() {
|
||||
AgentAguiRunJournal journal = new AgentAguiRunJournal(
|
||||
mock(AgentAguiRunStore.class), new ObjectMapper());
|
||||
String oversized = "x".repeat(512 * 1024 + 1);
|
||||
|
||||
Assert.assertThrows(IllegalStateException.class,
|
||||
() -> journal.append("run-1", oversized));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,124 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.springframework.data.redis.core.HashOperations;
|
||||
import org.springframework.data.redis.core.ListOperations;
|
||||
import org.springframework.data.redis.core.StringRedisTemplate;
|
||||
import org.springframework.data.redis.core.ValueOperations;
|
||||
import org.springframework.data.redis.core.script.RedisScript;
|
||||
import tech.easyflow.agent.config.AgentRuntimeProperties;
|
||||
import tech.easyflow.common.web.exceptions.BusinessException;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.ArgumentMatchers.anyList;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* AG-UI Redis 运行存储测试。
|
||||
*/
|
||||
public class AgentAguiRunStoreTest {
|
||||
|
||||
/**
|
||||
* 验证 runId 原子保留、24 小时 TTL 与完成终态写入。
|
||||
*/
|
||||
@Test
|
||||
public void shouldReserveRunAndPersistTerminalStatusForTwentyFourHours() {
|
||||
RedisMocks redis = redisMocks();
|
||||
AgentRuntimeProperties properties = new AgentRuntimeProperties();
|
||||
properties.setAguiRunRetention(Duration.ofHours(24));
|
||||
AgentAguiRunStore store = new AgentAguiRunStore(
|
||||
redis.template, new ObjectMapper(), properties);
|
||||
long now = System.currentTimeMillis();
|
||||
store.create(new AgentAguiRunDescriptor(
|
||||
"run-1", "request-1", "thread-1", "agent-1", "session-1",
|
||||
"user-1", "tenant-1", true, AgentAguiRunStatus.RUNNING, 0L, now, now));
|
||||
|
||||
when(redis.template.execute(any(RedisScript.class), anyList(), any(Object[].class)))
|
||||
.thenReturn(2L);
|
||||
store.append("run-1", "batch-1", List.of(
|
||||
"{\"type\":\"RUN_STARTED\",\"runId\":\"run-1\",\"threadId\":\"thread-1\"}",
|
||||
"{\"type\":\"RUN_FINISHED\",\"runId\":\"run-1\",\"threadId\":\"thread-1\"}"));
|
||||
|
||||
verify(redis.values).setIfAbsent(
|
||||
eq("easyflow:agent:agui:run:{run-1}:reserved"), eq("user-1"), eq(Duration.ofHours(24)));
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<Map<String, String>> metadata = ArgumentCaptor.forClass(Map.class);
|
||||
verify(redis.hashes).putAll(eq("easyflow:agent:agui:run:{run-1}:meta"), metadata.capture());
|
||||
Assert.assertEquals("0", metadata.getValue().get("lastCursor"));
|
||||
Assert.assertEquals("false", metadata.getValue().get("cancelRequested"));
|
||||
@SuppressWarnings("unchecked")
|
||||
ArgumentCaptor<RedisScript<Long>> script = ArgumentCaptor.forClass(RedisScript.class);
|
||||
ArgumentCaptor<Object[]> arguments = ArgumentCaptor.forClass(Object[].class);
|
||||
verify(redis.template).execute(script.capture(), anyList(), arguments.capture());
|
||||
Assert.assertTrue(script.getValue().getScriptAsString().contains("lastBatchId"));
|
||||
Assert.assertEquals("batch-1", arguments.getValue()[0]);
|
||||
Assert.assertEquals(AgentAguiRunStatus.COMPLETED.name(), arguments.getValue()[3]);
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证重复 runId 在模型启动前即被拒绝。
|
||||
*/
|
||||
@Test(expected = BusinessException.class)
|
||||
public void shouldRejectDuplicateRunId() {
|
||||
RedisMocks redis = redisMocks();
|
||||
when(redis.values.setIfAbsent(anyString(), anyString(), any(Duration.class))).thenReturn(false);
|
||||
AgentAguiRunStore store = new AgentAguiRunStore(
|
||||
redis.template, new ObjectMapper(), new AgentRuntimeProperties());
|
||||
store.create(new AgentAguiRunDescriptor(
|
||||
"run-1", "request-1", "thread-1", "agent-1", "session-1",
|
||||
"user-1", "tenant-1", false, AgentAguiRunStatus.RUNNING, 0L, 1L, 1L));
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证 owner 丢失通过同一原子批次追加明确失败终态。
|
||||
*/
|
||||
@Test
|
||||
public void shouldPersistOwnerLostAsReplayableTerminalEvent() {
|
||||
RedisMocks redis = redisMocks();
|
||||
when(redis.template.execute(any(RedisScript.class), anyList(), any(Object[].class)))
|
||||
.thenReturn(1L);
|
||||
AgentAguiRunStore store = new AgentAguiRunStore(
|
||||
redis.template, new ObjectMapper(), new AgentRuntimeProperties());
|
||||
AgentAguiRunDescriptor descriptor = new AgentAguiRunDescriptor(
|
||||
"run-1", "request-1", "thread-1", "agent-1", "session-1",
|
||||
"user-1", "tenant-1", false, AgentAguiRunStatus.RUNNING, 0L, 1L, 1L);
|
||||
ArgumentCaptor<Object[]> arguments = ArgumentCaptor.forClass(Object[].class);
|
||||
|
||||
store.failOwnerLost(descriptor);
|
||||
|
||||
verify(redis.template).execute(any(RedisScript.class), anyList(), arguments.capture());
|
||||
Assert.assertEquals(AgentAguiRunStatus.FAILED.name(), arguments.getValue()[3]);
|
||||
Assert.assertTrue(arguments.getValue()[8].toString().contains("AGENT_RUN_OWNER_LOST"));
|
||||
}
|
||||
|
||||
@SuppressWarnings("unchecked")
|
||||
private RedisMocks redisMocks() {
|
||||
StringRedisTemplate template = mock(StringRedisTemplate.class);
|
||||
ValueOperations<String, String> values = mock(ValueOperations.class);
|
||||
HashOperations<String, Object, Object> hashes = mock(HashOperations.class);
|
||||
ListOperations<String, String> lists = mock(ListOperations.class);
|
||||
when(template.opsForValue()).thenReturn(values);
|
||||
when(template.opsForHash()).thenReturn(hashes);
|
||||
when(template.opsForList()).thenReturn(lists);
|
||||
when(values.setIfAbsent(anyString(), anyString(), any(Duration.class))).thenReturn(true);
|
||||
return new RedisMocks(template, values, hashes, lists);
|
||||
}
|
||||
|
||||
private record RedisMocks(
|
||||
StringRedisTemplate template,
|
||||
ValueOperations<String, String> values,
|
||||
HashOperations<String, Object, Object> hashes,
|
||||
ListOperations<String, String> lists) {
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import org.junit.Test;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
|
||||
import java.util.List;
|
||||
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
/**
|
||||
* AG-UI 重连订阅共享调度测试。
|
||||
*/
|
||||
public class AgentAguiRunSubscriptionServiceTest {
|
||||
|
||||
/**
|
||||
* 验证同一 run 的不同游标独立读取窗口，尾部实时订阅不受历史回放阻塞。
|
||||
*
|
||||
* @throws Exception SSE mock 发送异常
|
||||
*/
|
||||
@Test
|
||||
public void shouldDispatchDistinctCursorWindowsIndependently() throws Exception {
|
||||
AgentAguiRunStore store = mock(AgentAguiRunStore.class);
|
||||
AgentAguiRunSubscriptionService service = new AgentAguiRunSubscriptionService(store);
|
||||
AgentAguiRunDescriptor descriptor = new AgentAguiRunDescriptor(
|
||||
"run-1", "request-1", "thread-1", "agent-1", "session-1",
|
||||
"user-1", "tenant-1", false, AgentAguiRunStatus.RUNNING,
|
||||
1000L, 1L, 2L);
|
||||
SseEmitter fromStart = mock(SseEmitter.class);
|
||||
SseEmitter fromTail = mock(SseEmitter.class);
|
||||
service.subscribe(descriptor, 0L, fromStart);
|
||||
service.subscribe(descriptor, 999L, fromTail);
|
||||
when(store.readAfter("run-1", 0L)).thenReturn(java.util.Collections.nCopies(
|
||||
256, "{\"type\":\"TEXT_MESSAGE_CONTENT\"}"));
|
||||
when(store.readAfter("run-1", 999L)).thenReturn(List.of(
|
||||
"{\"type\":\"TEXT_MESSAGE_CONTENT\",\"delta\":\"tail\"}"));
|
||||
when(store.find("run-1")).thenReturn(descriptor);
|
||||
|
||||
service.dispatch();
|
||||
|
||||
verify(store).readAfter("run-1", 0L);
|
||||
verify(store).readAfter("run-1", 999L);
|
||||
verify(store).find("run-1");
|
||||
verify(fromStart, times(256)).send(any(SseEmitter.SseEventBuilder.class));
|
||||
verify(fromTail).send(any(SseEmitter.SseEventBuilder.class));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package tech.easyflow.agent.runtime.agui;
|
||||
|
||||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.doThrow;
|
||||
import static org.mockito.Mockito.verify;
|
||||
|
||||
/**
|
||||
* 可重连 AG-UI SSE 发射器测试。
|
||||
*/
|
||||
public class ResumableAguiSseEmitterTest {
|
||||
|
||||
/**
|
||||
* 验证协议事件先进入日志，且逻辑关闭只发生在运行终态。
|
||||
*/
|
||||
@Test
|
||||
public void shouldKeepLogicalRunOpenUntilExplicitCompletion() {
|
||||
AgentAguiRunJournal journal = mock(AgentAguiRunJournal.class);
|
||||
SseEmitter subscriber = new SseEmitter();
|
||||
ResumableAguiSseEmitter emitter = new ResumableAguiSseEmitter(
|
||||
"thread-1", "run-1", journal, subscriber);
|
||||
|
||||
Assert.assertSame(subscriber, emitter.getEmitter());
|
||||
Assert.assertTrue(emitter.sendData("{\"type\":\"RUN_STARTED\"}"));
|
||||
Assert.assertFalse(emitter.isClosed());
|
||||
verify(journal).append("run-1", "{\"type\":\"RUN_STARTED\"}");
|
||||
|
||||
emitter.complete();
|
||||
|
||||
Assert.assertTrue(emitter.isClosed());
|
||||
Assert.assertTrue(emitter.isJournalCompletionSuccessful());
|
||||
verify(journal).complete("run-1");
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证日志拒绝普通事件时立即写入可重放故障终态并关闭逻辑输出。
|
||||
*/
|
||||
@Test
|
||||
public void shouldFailRunWhenJournalRejectsEvent() {
|
||||
AgentAguiRunJournal journal = mock(AgentAguiRunJournal.class);
|
||||
doThrow(new IllegalStateException("超过 16 MiB 安全上限"))
|
||||
.when(journal).append("run-1", "{\"type\":\"RUN_STARTED\"}");
|
||||
ResumableAguiSseEmitter emitter = new ResumableAguiSseEmitter(
|
||||
"thread-1", "run-1", journal, new SseEmitter());
|
||||
|
||||
Assert.assertFalse(emitter.sendData("{\"type\":\"RUN_STARTED\"}"));
|
||||
|
||||
Assert.assertTrue(emitter.isClosed());
|
||||
Assert.assertFalse(emitter.isJournalCompletionSuccessful());
|
||||
verify(journal).fail(
|
||||
"thread-1", "run-1", "AGENT_RUN_JOURNAL_LIMIT",
|
||||
"Agent 输出超过可恢复日志容量，运行已停止");
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证终态刷盘失败会被调用方观察到，避免上层错误记录为完成态。
|
||||
*/
|
||||
@Test
|
||||
public void shouldExposeTerminalJournalFailure() {
|
||||
AgentAguiRunJournal journal = mock(AgentAguiRunJournal.class);
|
||||
doThrow(new IllegalStateException("Redis unavailable"))
|
||||
.when(journal).complete("run-1");
|
||||
ResumableAguiSseEmitter emitter = new ResumableAguiSseEmitter(
|
||||
"thread-1", "run-1", journal, new SseEmitter());
|
||||
|
||||
emitter.complete();
|
||||
|
||||
Assert.assertTrue(emitter.isClosed());
|
||||
Assert.assertFalse(emitter.isJournalCompletionSuccessful());
|
||||
verify(journal).fail(
|
||||
"thread-1", "run-1", "AGENT_RUN_JOURNAL_UNAVAILABLE",
|
||||
"Agent 输出日志暂时不可用，运行已停止");
|
||||
}
|
||||
}
|
||||
@@ -6,6 +6,8 @@ import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
import org.mockito.ArgumentCaptor;
|
||||
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
|
||||
import tech.easyflow.agent.runtime.agui.AgentAguiRunJournal;
|
||||
import tech.easyflow.agent.runtime.agui.ResumableAguiSseEmitter;
|
||||
import tech.easyflow.core.chat.protocol.ChatDomain;
|
||||
import tech.easyflow.core.chat.protocol.ChatType;
|
||||
import tech.easyflow.core.chat.protocol.sse.ChatSseEmitter;
|
||||
@@ -51,6 +53,26 @@ public class AguiAgentRunOutputTest {
|
||||
verify(emitter).complete();
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证可恢复日志终态刷盘失败会沿输出边界返回失败。
|
||||
*/
|
||||
@Test
|
||||
public void shouldReportResumableJournalCompletionFailure() {
|
||||
AgentAguiRunJournal journal = mock(AgentAguiRunJournal.class);
|
||||
doThrow(new IllegalStateException("Redis unavailable"))
|
||||
.when(journal).complete("run-1");
|
||||
ResumableAguiSseEmitter emitter = new ResumableAguiSseEmitter(
|
||||
"123", "run-1", journal, new SseEmitter());
|
||||
AguiAgentRunOutput output = new AguiAgentRunOutput(
|
||||
"123", "run-1", "user-message-1", emitter);
|
||||
|
||||
Assert.assertTrue(output.emitRuntimeEvent(
|
||||
AgentRuntimeEvent.of(AgentRuntimeEventType.COMPLETED)));
|
||||
|
||||
Assert.assertFalse(output.finish("final"));
|
||||
Assert.assertFalse(emitter.isJournalCompletionSuccessful());
|
||||
}
|
||||
|
||||
/**
|
||||
* 验证审批 Custom Event 不泄漏恢复令牌。
|
||||
*/
|
||||
|
||||
Reference in New Issue
Block a user