Compare commits

6 Commits

Author SHA1 Message Date
9667a6d262 fix: 统一工作流失败原因与执行状态
- 关联 EASY-2,补齐模型错误分类、节点归属与执行日志

- 覆盖流式失败、重试状态及模型错误映射回归
2026-09-08 10:47:54 +08:00
130423edb4 feat: 支持工作流多知识库向量检索 2026-09-04 17:40:15 +08:00
7900552ede feat: 重构用户确认节点选择与恢复契约
- 统一单选多选输出与结构化暂停参数

- 增加严格恢复校验及并发状态保护

- 补充确认节点契约与恢复测试
2026-09-04 14:48:48 +08:00
368b90b211 perf: 复用 Milvus 客户端连接池 2026-09-04 11:22:13 +08:00
1ea7b7527f fix: 暴露本地文档解析任务丢失状态 2026-09-04 11:21:46 +08:00
68fd303656 fix: 拦截非标准 XLSX 文件
- 在 POI 解析前识别旧版 XLS 与异常容器

- 返回可操作提示并补充回归测试
2026-09-02 19:15:00 +08:00
48 changed files with 4861 additions and 384 deletions

View File

@@ -20,6 +20,7 @@ import com.easyagents.core.model.chat.ChatContext;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.model.exception.ModelException;
import com.easyagents.core.parser.AiMessageParser;
import com.easyagents.core.util.StringUtil;
import com.alibaba.fastjson2.JSON;
@@ -57,6 +58,7 @@ public class BaseStreamClientListener implements StreamClientListener {
@Override
public void onMessage(StreamClient client, String response) {
if (isFailure.get() || stoppedFlag.get()) return;
if (StringUtil.noText(response) || "[DONE]".equalsIgnoreCase(response.trim()) || finishedFlag.get()) {
notifyLastMessageAndStop(response);
return;
@@ -64,6 +66,12 @@ public class BaseStreamClientListener implements StreamClientListener {
try {
JSONObject jsonObject = JSON.parseObject(response);
if (jsonObject != null && jsonObject.get("error") != null) {
JSONObject error = jsonObject.get("error") instanceof JSONObject value ? value : new JSONObject();
String code = error.getString("code");
throw new ModelException(200, "Model stream error: " + response, null,
code, error.getString("type"), error.getString("message"));
}
AiMessage delta = messageParser.parse(jsonObject, chatContext);
//合并 增量 delta 到 fullMessage
@@ -77,8 +85,15 @@ public class BaseStreamClientListener implements StreamClientListener {
AiMessageResponse resp = new AiMessageResponse(chatContext, response, delta);
streamResponseListener.onMessage(context, resp);
} catch (Exception err) {
onFailure(this.context.getClient(), err);
onStop(this.context.getClient());
try {
onFailure(client, err);
} finally {
try {
client.stop();
} finally {
onStop(client);
}
}
}
}
@@ -133,6 +148,7 @@ public class BaseStreamClientListener implements StreamClientListener {
@Override
public void onFailure(StreamClient client, Throwable throwable) {
if (stoppedFlag.get() || finishedFlag.get()) return;
if (isFailure.compareAndSet(false, true)) {
context.setThrowable(throwable);
streamResponseListener.onFailure(context, throwable);

View File

@@ -15,6 +15,8 @@
*/
package com.easyagents.core.model.client.impl;
import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;
import com.easyagents.core.model.exception.ModelException;
import com.easyagents.core.util.StringUtil;
import okhttp3.Response;
@@ -25,11 +27,14 @@ import java.io.IOException;
class Util {
public static Throwable getFailureThrowable(Throwable t, Response response) {
if (t != null) {
if (t != null && response == null) {
return t;
}
if (response != null) {
String errorCode = null;
String errorType = null;
String errorMessage = null;
String errMessage = "Response code: " + response.code();
String message = response.message();
if (StringUtil.hasText(message)) {
@@ -40,12 +45,22 @@ class Util {
String string = body.string();
if (StringUtil.hasText(string)) {
errMessage += ", body: " + string;
try {
JSONObject payload = JSON.parseObject(string);
if (payload != null && payload.get("error") instanceof JSONObject error) {
errorCode = error.getString("code");
errorType = error.getString("type");
errorMessage = error.getString("message");
}
} catch (RuntimeException ignored) {
// 网关可能返回 HTML 或不完整 JSON仍保留 HTTP 状态和原始诊断。
}
}
}
} catch (IOException e) {
// ignore
}
t = new ModelException(errMessage);
t = new ModelException(response.code(), errMessage, t, errorCode, errorType, errorMessage);
}
return t;

View File

@@ -16,6 +16,38 @@
package com.easyagents.core.model.exception;
public class ModelException extends RuntimeException {
private Integer statusCode;
private String errorCode;
private String errorType;
private String errorMessage;
public ModelException(int statusCode, String message, Throwable cause) {
this(statusCode, message, cause, null, null, null);
}
public ModelException(int statusCode, String message, Throwable cause, String errorCode, String errorType, String errorMessage) {
super(message, cause);
this.statusCode = statusCode;
this.errorCode = errorCode;
this.errorType = errorType;
this.errorMessage = errorMessage;
}
public Integer getStatusCode() {
return statusCode;
}
public String getErrorCode() {
return errorCode;
}
public String getErrorMessage() {
return errorMessage;
}
public String getErrorType() {
return errorType;
}
/**
* Constructs a new runtime exception with {@code null} as its

View File

@@ -43,6 +43,11 @@ public class StoreOptions extends Metadata {
public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) {
throw new IllegalStateException("Can not set embeddingOptions to the default instance.");
}
@Override
public void setTimeoutMillis(Long timeoutMillis) {
throw new IllegalStateException("Can not set timeoutMillis to the default instance.");
}
};
/**
@@ -65,6 +70,11 @@ public class StoreOptions extends Metadata {
*/
private EmbeddingOptions embeddingOptions = EmbeddingOptions.DEFAULT;
/**
* Optional upper bound for one store operation.
*/
private Long timeoutMillis;
public String getCollectionName() {
return collectionName;
@@ -111,6 +121,17 @@ public class StoreOptions extends Metadata {
this.embeddingOptions = embeddingOptions;
}
public Long getTimeoutMillis() {
return timeoutMillis;
}
public void setTimeoutMillis(Long timeoutMillis) {
if (timeoutMillis != null && timeoutMillis <= 0L) {
throw new IllegalArgumentException("timeoutMillis must be greater than zero");
}
this.timeoutMillis = timeoutMillis;
}
public static StoreOptions ofCollectionName(String collectionName) {
StoreOptions storeOptions = new StoreOptions();

View File

@@ -0,0 +1,21 @@
/*
* Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
*/
package com.easyagents.core.store;
/**
* Indicates that a store operation exhausted its caller-provided time budget.
*/
public class StoreTimeoutException extends RuntimeException {
public StoreTimeoutException(String message) {
super(message);
}
public StoreTimeoutException(String message, Throwable cause) {
super(message, cause);
}
}

View File

@@ -0,0 +1,52 @@
package com.easyagents.core.model.client;
import com.easyagents.core.message.AiMessage;
import com.easyagents.core.model.chat.ChatConfig;
import com.easyagents.core.model.chat.ChatContext;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import org.junit.Assert;
import org.junit.Test;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
public class StreamFailureLifecycleTest {
@Test
public void completedStreamMustIgnoreLateFailureAndData() {
Fixture fixture = new Fixture();
fixture.listener.onMessage(fixture.client, "{}");
fixture.listener.onMessage(fixture.client, "[DONE]");
fixture.listener.onFailure(fixture.client, new IllegalStateException("late transport error"));
fixture.listener.onMessage(fixture.client, "{}");
Assert.assertEquals(List.of("message", "message", "stop"), fixture.events);
}
@Test
public void failedStreamMustStopWithoutLaterSuccess() {
Fixture fixture = new Fixture();
fixture.listener.onMessage(fixture.client, "{}");
fixture.listener.onMessage(fixture.client, "{\"error\":{\"code\":\"rate_limit_exceeded\"}}");
fixture.listener.onMessage(fixture.client, "[DONE]");
fixture.listener.onMessage(fixture.client, "{}");
fixture.listener.onFailure(fixture.client, new IllegalStateException("cancelled"));
Assert.assertEquals(List.of("message", "failure", "transport-stop", "stop"), fixture.events);
}
private static class Fixture {
final List<String> events = new ArrayList<>();
final StreamClient client = new StreamClient() {
public void start(String url, Map<String, String> headers, String payload, StreamClientListener listener, ChatConfig config) { }
public void stop() { events.add("transport-stop"); }
};
final BaseStreamClientListener listener = new BaseStreamClientListener(null, new ChatContext(), client,
new StreamResponseListener() {
public void onMessage(StreamContext context, AiMessageResponse response) { events.add("message"); }
public void onFailure(StreamContext context, Throwable error) { events.add("failure"); }
public void onStop(StreamContext context) { events.add("stop"); }
}, (json, context) -> {
AiMessage message = new AiMessage(); message.setContent("partial"); return message;
});
}
}

View File

@@ -0,0 +1,39 @@
package com.easyagents.core.model.client.impl;
import com.easyagents.core.model.exception.ModelException;
import okhttp3.MediaType;
import okhttp3.Protocol;
import okhttp3.Request;
import okhttp3.Response;
import okhttp3.ResponseBody;
import org.junit.Assert;
import org.junit.Test;
public class ModelHttpFailureTest {
@Test
public void shouldPreserveStructuredErrorFromActualProviderShape() {
ModelException error = failure("{\"error\":{\"message\":\"Model not found\",\"code\":404,\"type\":\"NotFound\"}}");
Assert.assertEquals(Integer.valueOf(404), error.getStatusCode());
Assert.assertEquals("404", error.getErrorCode());
Assert.assertEquals("NotFound", error.getErrorType());
Assert.assertEquals("Model not found", error.getErrorMessage());
Assert.assertTrue(error.getMessage().contains("Model not found"));
}
@Test
public void shouldPreserveHttpStatusForNonModelGatewayErrors() {
for (String body : new String[]{"<html>Not Found</html>", "{broken", "{\"error\":\"not found\"}", "null"}) {
ModelException error = failure(body);
Assert.assertEquals(Integer.valueOf(404), error.getStatusCode());
Assert.assertNull(error.getErrorCode());
Assert.assertNull(error.getErrorMessage());
}
}
private ModelException failure(String body) {
Response response = new Response.Builder().request(new Request.Builder().url("http://localhost/chat").build())
.protocol(Protocol.HTTP_1_1).code(404).message("Not Found")
.body(ResponseBody.create(MediaType.get("application/json"), body)).build();
return (ModelException) Util.getFailureThrowable(null, response);
}
}

View File

@@ -2,6 +2,7 @@ package com.easyagents.document.core.async;
import com.easyagents.core.util.StringUtil;
import com.easyagents.document.core.exception.DocumentParseException;
import com.easyagents.document.core.exception.DocumentAsyncTaskNotFoundException;
import com.easyagents.document.core.entity.ParseResponse;
import com.easyagents.document.core.entity.ParseTaskInfo;
import com.easyagents.document.core.entity.ParseTaskStatus;
@@ -135,7 +136,7 @@ public class DocumentAsyncTaskManager {
}
DocumentAsyncTaskRecord record = repository.find(taskId);
if (record == null) {
throw new DocumentParseException("Document async task not found: " + taskId);
throw new DocumentAsyncTaskNotFoundException(taskId);
}
return record;
}

View File

@@ -0,0 +1,29 @@
package com.easyagents.document.core.exception;
/**
* 进程内异步文档任务不存在。
*
* <p>本地 Office 解析任务允许使用内存仓库;进程重启后,调用方可以
* 通过该异常识别执行实例已经丢失,并从持久化业务任务重新提交。</p>
*
* @author Codex
* @since 2026-09-02
*/
public class DocumentAsyncTaskNotFoundException extends DocumentParseException {
private final String taskId;
public DocumentAsyncTaskNotFoundException(String taskId) {
super("Document async task not found: " + taskId);
this.taskId = taskId;
}
/**
* 获取已丢失的任务 ID。
*
* @return 任务 ID
*/
public String getTaskId() {
return taskId;
}
}

View File

@@ -4,6 +4,7 @@ import com.easyagents.document.core.entity.ParseResponse;
import com.easyagents.document.core.entity.ParseResult;
import com.easyagents.document.core.entity.ParseTaskInfo;
import com.easyagents.document.core.entity.ParseTaskStatus;
import com.easyagents.document.core.exception.DocumentAsyncTaskNotFoundException;
import org.junit.Assert;
import org.junit.Test;
@@ -18,6 +19,21 @@ import java.util.concurrent.Executor;
*/
public class DocumentAsyncTaskManagerTest {
@Test
public void shouldExposeMissingInMemoryTaskAsRecoverableSignal() {
DocumentAsyncTaskManager manager = new DocumentAsyncTaskManager(
new InMemoryDocumentAsyncTaskRepository(),
Runnable::run
);
try {
manager.queryTaskInfo("lost-task");
Assert.fail("expected DocumentAsyncTaskNotFoundException");
} catch (DocumentAsyncTaskNotFoundException error) {
Assert.assertEquals("lost-task", error.getTaskId());
}
}
@Test
public void shouldTrackTaskLifecycleAndResult() {
Executor directExecutor = new Executor() {

View File

@@ -14,6 +14,7 @@ import com.easyagents.document.core.entity.ParseRequest;
import com.easyagents.document.core.entity.ParseResponse;
import com.easyagents.document.core.entity.ParseResult;
import com.easyagents.document.core.entity.XlsxParseRequest;
import com.easyagents.document.core.exception.DocumentParseException;
import com.easyagents.document.core.support.AbstractAsyncDocumentParseService;
import com.easyagents.document.xlsx.XlsxDocumentProvider;
import com.easyagents.document.xlsx.model.XlsxCellArtifact;
@@ -52,6 +53,10 @@ import java.util.concurrent.Executors;
public class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseService<XlsxParseRequest> implements XlsxDocumentProvider {
public static final String PROVIDER_NAME = "mineru";
private static final byte[] OLE2_SIGNATURE = new byte[] {
(byte) 0xd0, (byte) 0xcf, 0x11, (byte) 0xe0,
(byte) 0xa1, (byte) 0xb1, 0x1a, (byte) 0xe1
};
private final MineruProperties properties;
private final MineruClient client;
@@ -145,6 +150,9 @@ public class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseSe
@Override
protected ParseResponse doParse(XlsxParseRequest request, DocumentAsyncTaskUpdater updater) {
for (ParseFile file : request.getFiles()) {
validateXlsxContent(file);
}
ParseResponse response = new ParseResponse();
List<ParseResult> results = new ArrayList<ParseResult>();
String backend = null;
@@ -211,6 +219,50 @@ public class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseSe
return aggregate;
}
/**
* 在 POI 打开工作簿前校验 XLSX 容器签名,避免向用户暴露底层格式异常。
*
* @param file 待解析文件
*/
private void validateXlsxContent(ParseFile file) {
byte[] content = file == null ? null : file.getContent();
if (hasZipSignature(content)) {
return;
}
String fileName = file == null || !StringUtil.hasText(file.getFileName())
? "当前文件"
: "文件“" + file.getFileName() + "";
String reason = startsWith(content, OLE2_SIGNATURE)
? "可能是旧版 XLS 或已加密文件"
: "文件内容与 .xlsx 扩展名不一致或文件已损坏";
throw new DocumentParseException(
fileName + "不是标准 XLSX" + reason
+ "。请解除保护后用 Excel/WPS 另存为 XLSX修改文件后缀无效"
);
}
private boolean hasZipSignature(byte[] content) {
return content != null
&& content.length >= 4
&& content[0] == 'P'
&& content[1] == 'K'
&& ((content[2] == 3 && content[3] == 4)
|| (content[2] == 5 && content[3] == 6)
|| (content[2] == 7 && content[3] == 8));
}
private boolean startsWith(byte[] content, byte[] signature) {
if (content == null || content.length < signature.length) {
return false;
}
for (int index = 0; index < signature.length; index++) {
if (content[index] != signature[index]) {
return false;
}
}
return true;
}
private SheetExtraction extractSheet(XSSFSheet sheet,
int sheetIndex,
DataFormatter formatter,

View File

@@ -16,6 +16,7 @@ import com.easyagents.document.core.entity.ParseTaskStatus;
import com.easyagents.document.core.entity.XlsxParseRequest;
import com.easyagents.document.core.exception.DocumentParseException;
import com.easyagents.document.xlsx.model.XlsxParseArtifact;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.ss.usermodel.ClientAnchor;
import org.apache.poi.xssf.usermodel.XSSFDrawing;
import org.apache.poi.xssf.usermodel.XSSFSheet;
@@ -119,6 +120,31 @@ public class MineruXlsxDocumentParseServiceTest {
Assert.assertEquals("image/jpeg", result.getImages().get(0).getMimeType());
}
@Test
public void shouldRejectLegacyXlsContentWithActionableMessage() throws Exception {
RecordingClient client = new RecordingClient(defaultProperties());
MineruMapper mapper = new MineruMapper(defaultProperties());
MineruXlsxDocumentParseService service = new MineruXlsxDocumentParseService(
defaultProperties(),
client,
mapper,
new DocumentAsyncTaskManager(new InMemoryDocumentAsyncTaskRepository(), directExecutor())
);
XlsxParseRequest request = new XlsxParseRequest();
request.addFile(ParseFile.of("legacy.xlsx", buildLegacyWorkbookBytes()));
DocumentParseException error = Assert.assertThrows(
DocumentParseException.class,
() -> service.parse(request)
);
Assert.assertEquals(
"文件“legacy.xlsx”不是标准 XLSX可能是旧版 XLS 或已加密文件。"
+ "请解除保护后用 Excel/WPS 另存为 XLSX修改文件后缀无效",
error.getMessage()
);
}
@Test
public void shouldAppendImageReferenceForImageOnlySheet() throws Exception {
RecordingClient client = new RecordingClient(defaultProperties());
@@ -255,6 +281,15 @@ public class MineruXlsxDocumentParseServiceTest {
return writeWorkbook(workbook);
}
private byte[] buildLegacyWorkbookBytes() throws Exception {
try (HSSFWorkbook workbook = new HSSFWorkbook();
ByteArrayOutputStream outputStream = new ByteArrayOutputStream()) {
workbook.createSheet("Sheet1").createRow(0).createCell(0).setCellValue("旧版表格");
workbook.write(outputStream);
return outputStream.toByteArray();
}
}
private void addPicture(XSSFWorkbook workbook,
XSSFSheet sheet,
int rowIndex,

View File

@@ -612,6 +612,8 @@ public class Chain {
Map<String, Object> nodeResult = null;
Throwable error = null;
String executionAttemptKey = null;
String auditInstanceId = StringUtil.hasText(chainState.getAuditInstanceId())
? chainState.getAuditInstanceId() : stateInstanceId;
try {
AtomicBoolean nodeStarted = new AtomicBoolean();
String candidateAttemptKey =
@@ -647,12 +649,13 @@ public class Chain {
activeNodeState
.getExecutionAttemptKey();
if (nodeStarted.get()) {
auditInstanceId = getAuditInstanceId();
notifyEvent(new NodeStartEvent(
this,
node,
executionAttemptKey,
activeNodeState.getStatus(),
getAuditInstanceId()));
auditInstanceId));
}
ChainState nodeExecutionState = updateStateSafely(state -> {
@@ -672,7 +675,10 @@ public class Chain {
} catch (TriggerClaimLostException | RetryableTriggerException claimLost) {
throw claimLost;
} catch (Throwable throwable) {
log.error("Node execute error", throwable);
if (!(throwable instanceof ChainSuspendException)) {
log.error("Node execute error, executeId={}, chainInstanceId={}, nodeId={}, nodeName={}, attemptKey={}",
auditInstanceId, stateInstanceId, node.getId(), node.getName(), executionAttemptKey, throwable);
}
error = throwable;
}
// 结果提交入口会在实例锁内重读状态,统一拦截取消、超时和其他终态。
@@ -882,6 +888,10 @@ public class Chain {
NodeStatus finalNodeStatus = null;
try {
if (error == null) {
updateNodeStateSafely(node.id, state -> {
state.setError(null);
return EnumSet.of(NodeStateField.ERROR);
});
// 更新 state 数据
updateStateSafely(state -> {
EnumSet<ChainStateField> fields = EnumSet.of(ChainStateField.EXECUTE_RESULT);
@@ -941,9 +951,10 @@ public class Chain {
}
// 失败
else {
finalNodeStatus = NodeStatus.ERROR;
NodeState newState = updateNodeStateSafely(node.getId(), s -> {
s.setStatus(NodeStatus.ERROR);
s.setError(new ExceptionSummary(error));
s.setError(new ExceptionSummary(error, stateInstanceId, node.getId(), node.getName()));
return EnumSet.of(NodeStateField.ERROR, NodeStateField.STATUS);
});
@@ -960,7 +971,8 @@ public class Chain {
scheduleNode(node, triggerEdgeId, TriggerType.RETRY, node.getRetryIntervalMs());
} else {
finalChainStatus = handleNodeError(node.id, error);
finalNodeStatus = NodeStatus.FAILED;
finalChainStatus = handleNodeError(node, error);
}
}
}
@@ -992,10 +1004,15 @@ public class Chain {
// 更新父级链的状态
if (!finalChainStatus.isSuccess()) {
ChainState currentState = getState();
ExceptionSummary originError = currentState.getError();
ChainStatus currentStatus = finalChainStatus;
while (currentState != null && StringUtil.hasText(currentState.getParentInstanceId())) {
updateStateSafely(currentState.getParentInstanceId(), state -> {
state.setStatus(currentStatus);
if (state.getError() == null && originError != null) {
state.setError(originError);
return EnumSet.of(ChainStateField.STATUS, ChainStateField.ERROR);
}
return EnumSet.of(ChainStateField.STATUS);
});
setStatusAndNotifyEvent(currentState.getParentInstanceId(), currentStatus);
@@ -1661,7 +1678,7 @@ public class Chain {
}
before.set(state.getStatus());
state.setStatus(ChainStatus.FAILED);
state.setError(new ExceptionSummary(failure));
state.setError(new ExceptionSummary(failure, stateInstanceId, null, null));
state.setMessage(failure.getMessage());
changed.set(true);
return EnumSet.of(
@@ -1693,19 +1710,23 @@ public class Chain {
}
private ChainStatus handleNodeError(String nodeId, Throwable throwable) {
updateNodeStateSafely(nodeId, s -> {
private ChainStatus handleNodeError(Node node, Throwable throwable) {
ExceptionSummary summary = new ExceptionSummary(throwable, stateInstanceId, node.getId(), node.getName());
updateNodeStateSafely(node.getId(), s -> {
s.setStatus(NodeStatus.FAILED);
s.setError(new ExceptionSummary(throwable));
s.setError(summary);
return EnumSet.of(NodeStateField.ERROR, NodeStateField.STATUS);
});
updateStateSafely(state -> {
state.setError(new ExceptionSummary(throwable));
if (state.getError() != null) {
return null;
}
state.setError(summary);
return EnumSet.of(ChainStateField.ERROR);
});
setStatusAndNotifyEvent(ChainStatus.FAILED);
// 节点结束事件先于工作流终态,终态订阅者关闭连接前能收到失败节点。
deferOrRun(() -> eventManager.notifyChainError(throwable, this));
return ChainStatus.FAILED;
}
@@ -1751,16 +1772,91 @@ public class Chain {
stateInstanceId,
10L,
TimeUnit.SECONDS,
() -> {
ChainState current =
chainStateRepository.load(stateInstanceId);
if (current == null
|| current.getStatus() != ChainStatus.SUSPEND) {
return false;
() -> resumeSuspended(variables));
}
resumeSuspended(variables);
private void validateResumeVariables(
ChainState state, Map<String, Object> variables) {
List<Parameter> parameters = state.getSuspendForParameters();
if (parameters == null || parameters.isEmpty()) {
return;
}
Map<String, Object> submitted = variables == null
? Collections.emptyMap()
: variables;
Set<String> expectedKeys = parameters.stream()
.map(Parameter::getName)
.filter(Objects::nonNull)
.collect(java.util.stream.Collectors.toCollection(
LinkedHashSet::new));
Set<String> extraKeys = new LinkedHashSet<>(submitted.keySet());
extraKeys.removeAll(expectedKeys);
if (!extraKeys.isEmpty()) {
throw new ChainResumeException(
"确认参数包含未声明字段");
}
for (Parameter parameter : parameters) {
String name = parameter.getName();
Object value = submitted.get(name);
String label = StringUtil.getFirstWithText(
parameter.getFormLabel(), name);
if (!submitted.containsKey(name) || isBlankResumeValue(value)) {
throw new ChainResumeException(
"确认参数[" + label + "]不能为空");
}
validateResumeOption(parameter, value, label);
}
}
private void validateResumeOption(
Parameter parameter, Object value, String label) {
List<ParameterOption> options = parameter.getOptions();
if (options == null || options.isEmpty()) {
return;
}
Set<String> allowedValues = options.stream()
.map(ParameterOption::getValue)
.filter(Objects::nonNull)
.collect(java.util.stream.Collectors.toCollection(
LinkedHashSet::new));
if ("checkbox".equals(parameter.getFormType())) {
if (!(value instanceof Collection<?> selected)) {
throw new ChainResumeException(
"确认参数[" + label + "]必须提交字符串数组");
}
Set<String> unique = new LinkedHashSet<>();
for (Object item : selected) {
if (!(item instanceof String selectedValue)
|| !allowedValues.contains(selectedValue)) {
throw new ChainResumeException(
"确认参数[" + label + "]包含未配置选项");
}
if (!unique.add(selectedValue)) {
throw new ChainResumeException(
"确认参数[" + label + "]不能重复选择同一选项");
}
}
return;
}
if (!(value instanceof String selectedValue)) {
throw new ChainResumeException(
"确认参数[" + label + "]必须提交单个字符串值");
}
if (!allowedValues.contains(selectedValue)) {
throw new ChainResumeException(
"确认参数[" + label + "]包含未配置选项");
}
}
private boolean isBlankResumeValue(Object value) {
if (value == null) {
return true;
});
}
if (value instanceof String text) {
return text.trim().isEmpty();
}
return value instanceof Collection<?> collection
&& collection.isEmpty();
}
/**
@@ -1777,36 +1873,51 @@ public class Chain {
*
* @param variables 恢复时注入的变量
*/
private void resumeSuspended(Map<String, Object> variables) {
ChainState newState = updateStateSafely(state -> {
if (variables != null) {
state.getMemory().putAll(variables);
return EnumSet.of(ChainStateField.MEMORY);
} else {
private boolean resumeSuspended(Map<String, Object> variables) {
AtomicBoolean resumed = new AtomicBoolean(false);
AtomicReference<Set<String>> suspendedNodeIds =
new AtomicReference<>(Collections.emptySet());
updateStateSafely(state -> {
resumed.set(false);
suspendedNodeIds.set(Collections.emptySet());
if (state.getStatus() != ChainStatus.SUSPEND) {
return null;
}
});
notifyEvent(new ChainResumeEvent(this, variables));
setStatusAndNotifyEvent(ChainStatus.RUNNING);
Set<String> suspendNodeIds = newState.getSuspendNodeIds();
if (suspendNodeIds != null && !suspendNodeIds.isEmpty()) {
// 移除 suspend 状态,方便二次 suspend 时,不带有旧数据
updateStateSafely(state -> {
validateResumeVariables(state, variables);
if (state.getSuspendNodeIds() != null) {
suspendedNodeIds.set(
new LinkedHashSet<>(state.getSuspendNodeIds()));
}
EnumSet<ChainStateField> updatedFields = EnumSet.of(
ChainStateField.STATUS,
ChainStateField.SUSPEND_NODE_IDS,
ChainStateField.SUSPEND_FOR_PARAMETERS);
if (variables != null && !variables.isEmpty()) {
state.getMemory().putAll(variables);
updatedFields.add(ChainStateField.MEMORY);
}
state.setStatus(ChainStatus.RUNNING);
state.setSuspendNodeIds(null);
state.setSuspendForParameters(null);
return EnumSet.of(ChainStateField.SUSPEND_NODE_IDS, ChainStateField.SUSPEND_FOR_PARAMETERS);
resumed.set(true);
return updatedFields;
});
if (!resumed.get()) {
return false;
}
for (String id : suspendNodeIds) {
notifyEvent(new ChainResumeEvent(this, variables));
notifyEvent(new ChainStatusChangeEvent(
this, ChainStatus.RUNNING, ChainStatus.SUSPEND));
for (String id : suspendedNodeIds.get()) {
Node node = definition.getNodeById(id);
if (node == null) {
throw new ChainException("Node not found: " + id);
}
scheduleNode(node, null, TriggerType.RESUME, 0L);
}
}
return true;
}
public void resume() {

View File

@@ -0,0 +1,16 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
*/
package com.easyagents.flow.core.chain;
/**
* 工作流挂起参数不满足当前恢复请求时抛出的异常。
*/
public class ChainResumeException extends ChainException {
public ChainResumeException(String message) {
super(message);
}
}

View File

@@ -596,6 +596,7 @@ public class ChainState implements Serializable {
List<Parameter> suspendParameters = null;
List<Map<String, Object>> templateRootMaps = null;
for (Parameter parameter : parameters) {
try {
RefType refType = parameter.getRefType();
Object value = null;
if (refType == RefType.FIXED) {
@@ -637,7 +638,8 @@ public class ChainState implements Serializable {
if (parameter.isRequired() && isNullOrBlank(value)) {
if (!ignoreRequired) {
throw new ChainException(node.getName() + " Missing required parameter:" + parameter.getName());
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
node.getName() + " Missing required parameter:" + parameter.getName());
}
}
@@ -653,6 +655,12 @@ public class ChainState implements Serializable {
}
variables.put(parameter.getName(), value);
} catch (WorkflowExecutionException failure) {
throw failure;
} catch (RuntimeException failure) {
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
"Cannot resolve parameter: " + parameter.getName(), failure);
}
}
if (suspendParameters != null && !suspendParameters.isEmpty()) {

View File

@@ -18,6 +18,9 @@ package com.easyagents.flow.core.chain;
import java.io.PrintWriter;
import java.io.Serializable;
import java.io.StringWriter;
import java.util.Collections;
import java.util.IdentityHashMap;
import java.util.Set;
public class ExceptionSummary implements Serializable {
private static final long serialVersionUID = 1L;
@@ -31,12 +34,20 @@ public class ExceptionSummary implements Serializable {
private String chainId;
private String nodeId;
private String nodeName;
private String errorCode; // 可选
private long timestamp;
public ExceptionSummary(Throwable error) {
this(error, null, null, null);
}
public ExceptionSummary(Throwable error, String chainId, String nodeId, String nodeName) {
this.chainId = chainId;
this.nodeId = nodeId;
this.nodeName = nodeName;
this.exceptionClass = error.getClass().getName();
this.message = error.getMessage();
this.stackTrace = getStackTraceAsString(error);
@@ -46,11 +57,30 @@ public class ExceptionSummary implements Serializable {
this.rootCauseMessage = root.getMessage();
this.timestamp = System.currentTimeMillis();
Set<Throwable> seen = Collections.newSetFromMap(new IdentityHashMap<>());
for (Throwable cause = error; cause != null && seen.add(cause); cause = cause.getCause()) {
if (cause instanceof WorkflowExecutionException failure) {
if (this.errorCode == null && failure.getReason() != null) {
this.errorCode = failure.getReason().getCode();
}
if (failure.getNodeId() != null) {
this.chainId = failure.getChainId();
this.nodeId = failure.getNodeId();
this.nodeName = failure.getNodeName();
}
}
}
if (this.errorCode == null) {
this.errorCode = (this.nodeId == null ? WorkflowErrorReason.WORKFLOW_INTERNAL_ERROR
: WorkflowErrorReason.NODE_EXECUTION_FAILED).getCode();
}
}
private static Throwable getRootCause(Throwable t) {
Throwable result = t;
while (result.getCause() != null) {
Set<Throwable> seen = Collections.newSetFromMap(new IdentityHashMap<>());
seen.add(result);
while (result.getCause() != null && seen.add(result.getCause())) {
result = result.getCause();
}
return result;
@@ -123,6 +153,10 @@ public class ExceptionSummary implements Serializable {
return errorCode;
}
public String getNodeName() { return nodeName; }
public void setNodeName(String nodeName) { this.nodeName = nodeName; }
public void setErrorCode(String errorCode) {
this.errorCode = errorCode;
}

View File

@@ -49,6 +49,11 @@ public class Parameter implements Serializable, Cloneable {
*/
protected List<Object> enums;
/**
* 显示文案与实际值分离的结构化选项。
*/
protected List<ParameterOption> options;
/**
* 用户输入的表单类型,例如:"input" "textarea" "select" "radio" "checkbox" 等等
*/
@@ -242,6 +247,14 @@ public class Parameter implements Serializable, Cloneable {
}
}
public List<ParameterOption> getOptions() {
return options;
}
public void setOptions(List<ParameterOption> options) {
this.options = options;
}
public String getFormType() {
return formType;
}
@@ -298,6 +311,7 @@ public class Parameter implements Serializable, Cloneable {
", flattenAggregation=" + flattenAggregation +
", children=" + children +
", enums=" + enums +
", options=" + options +
", formType='" + formType + '\'' +
", formLabel='" + formLabel + '\'' +
", formPlaceholder='" + formPlaceholder + '\'' +
@@ -320,6 +334,13 @@ public class Parameter implements Serializable, Cloneable {
clone.enums = new ArrayList<>(this.enums.size());
clone.enums.addAll(this.enums);
}
if (this.options != null) {
clone.options = new ArrayList<>(this.options.size());
for (ParameterOption option : this.options) {
clone.options.add(new ParameterOption(
option.getLabel(), option.getValue()));
}
}
return clone;
} catch (CloneNotSupportedException e) {
throw new AssertionError();

View File

@@ -0,0 +1,55 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
*/
package com.easyagents.flow.core.chain;
import java.io.Serializable;
/**
* 用户输入参数的结构化选项。
*/
public class ParameterOption implements Serializable {
private static final long serialVersionUID = 1L;
private String label;
private String value;
public ParameterOption() {
}
public ParameterOption(String label, String value) {
setLabel(label);
setValue(value);
}
public String getLabel() {
return label;
}
public void setLabel(String label) {
this.label = trim(label);
}
public String getValue() {
return value;
}
public void setValue(String value) {
this.value = trim(value);
}
private static String trim(String value) {
return value == null ? null : value.trim();
}
@Override
public String toString() {
return "ParameterOption{" +
"label='" + label + '\'' +
", value='" + value + '\'' +
'}';
}
}

View File

@@ -0,0 +1,37 @@
package com.easyagents.flow.core.chain;
/** 工作流对外稳定的失败原因;文案不包含底层响应、凭据或业务输入。 */
public enum WorkflowErrorReason {
INPUT_INVALID("输入参数无效,请检查必填参数和变量引用"),
MODEL_NOT_FOUND("模型不存在,请检查模型名称或重新选择模型"),
MODEL_RATE_LIMITED("模型调用受到限流,请稍后重试"),
MODEL_AUTH_FAILED("模型鉴权失败,请检查模型凭据和访问权限"),
MODEL_UNAVAILABLE("模型服务暂不可用,请稍后重试"),
MODEL_TIMEOUT("模型响应超时,请稍后重试"),
NODE_OUTPUT_INVALID("节点输出为空或格式不符合要求,请检查输出配置"),
NODE_EXECUTION_FAILED("节点执行失败"),
WORKFLOW_INTERNAL_ERROR("工作流内部执行异常");
private final String defaultMessage;
WorkflowErrorReason(String defaultMessage) {
this.defaultMessage = defaultMessage;
}
public String getCode() {
return name();
}
public String getDefaultMessage() {
return defaultMessage;
}
public static WorkflowErrorReason fromCode(String code) {
for (WorkflowErrorReason reason : values()) {
if (reason.getCode().equals(code)) {
return reason;
}
}
return null;
}
}

View File

@@ -0,0 +1,30 @@
package com.easyagents.flow.core.chain;
/** 保留原始 cause并在已知业务边界标注失败原因。 */
public class WorkflowExecutionException extends ChainException {
private final WorkflowErrorReason reason;
private String chainId;
private String nodeId;
private String nodeName;
public WorkflowExecutionException(WorkflowErrorReason reason, String message) {
this(reason, message, null);
}
public WorkflowExecutionException(WorkflowErrorReason reason, String message, Throwable cause) {
super(message, cause);
this.reason = reason;
}
public WorkflowExecutionException withContext(String chainId, String nodeId, String nodeName) {
this.chainId = chainId;
this.nodeId = nodeId;
this.nodeName = nodeName;
return this;
}
public WorkflowErrorReason getReason() { return reason; }
public String getChainId() { return chainId; }
public String getNodeId() { return nodeId; }
public String getNodeName() { return nodeName; }
}

View File

@@ -19,6 +19,8 @@ package com.easyagents.flow.core.chain.event;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.Node;
import com.easyagents.flow.core.chain.NodeStatus;
import com.easyagents.flow.core.chain.ExceptionSummary;
import com.easyagents.flow.core.chain.ChainSuspendException;
import java.util.Map;
@@ -30,6 +32,7 @@ public class NodeEndEvent extends BaseEvent {
private final Node node;
private final Map<String, Object> result;
private final Throwable error;
private final ExceptionSummary errorSummary;
private final NodeStatus status;
private final String executionAttemptKey;
@@ -88,6 +91,8 @@ public class NodeEndEvent extends BaseEvent {
this.node = node;
this.result = result;
this.error = error;
this.errorSummary = error == null || error instanceof ChainSuspendException ? null
: new ExceptionSummary(error, chain.getStateInstanceId(), node.getId(), node.getName());
this.status = status;
this.executionAttemptKey = executionAttemptKey;
}
@@ -119,6 +124,8 @@ public class NodeEndEvent extends BaseEvent {
return error;
}
public ExceptionSummary getErrorSummary() { return errorSummary; }
/**
* 获取事件创建时捕获的节点终态。
*

View File

@@ -898,6 +898,16 @@ public class ChainExecutor {
});
}
return node.execute(temp);
} catch (ChainSuspendException | TriggerClaimLostException | RetryableTriggerException control) {
throw control;
} catch (RuntimeException failure) {
log.error("Single node execute error, executeId={}, chainInstanceId={}, nodeId={}, nodeName={}, attemptKey={}",
temp.getStateInstanceId(), temp.getStateInstanceId(), nodeId, node.getName(),
temp.getStateInstanceId() + ":" + nodeId + ":single", failure);
ExceptionSummary summary = new ExceptionSummary(failure, temp.getStateInstanceId(), node.getId(), node.getName());
throw new WorkflowExecutionException(WorkflowErrorReason.fromCode(summary.getErrorCode()),
"Single node execution failed", failure)
.withContext(temp.getStateInstanceId(), summary.getNodeId(), summary.getNodeName());
} finally {
activeDefinitions.remove(temp.getStateInstanceId());
definitionSnapshotRepository.remove(temp.getStateInstanceId());

View File

@@ -18,6 +18,7 @@ package com.easyagents.flow.core.knowledge;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
public class KnowledgeManager {
@@ -51,4 +52,20 @@ public class KnowledgeManager {
}
return null;
}
/**
* 将完整知识检索请求交给首个能够处理它的 Provider。
*
* @param request 检索请求
* @return 节点输出;没有 Provider 能处理时返回 null
*/
public Map<String, Object> search(KnowledgeSearchRequest request) {
for (KnowledgeProvider provider : providers) {
Map<String, Object> result = provider.search(request);
if (result != null) {
return result;
}
}
return null;
}
}

View File

@@ -15,6 +15,37 @@
*/
package com.easyagents.flow.core.knowledge;
import com.easyagents.flow.core.util.Maps;
import java.util.List;
import java.util.Map;
public interface KnowledgeProvider {
Knowledge getKnowledge(Object id);
/**
* 执行完整的知识库节点检索请求。
*
* <p>默认实现保留单知识库兼容。需要跨知识库汇总的业务 Provider
* 应覆盖本方法并返回完整节点输出。</p>
*
* @param request 检索请求
* @return 节点输出;当前 Provider 不支持该请求时返回 null
*/
default Map<String, Object> search(KnowledgeSearchRequest request) {
if (request == null || request.getKnowledgeIds().size() != 1) {
return null;
}
Object knowledgeId = request.getKnowledgeIds().get(0);
Knowledge knowledge = getKnowledge(knowledgeId);
if (knowledge == null) {
return null;
}
List<Map<String, Object>> documents = knowledge.search(
request.getKeyword(),
request.getLimit(),
request.getKnowledgeNode(),
request.getChain());
return Maps.of("documents", documents);
}
}

View File

@@ -0,0 +1,67 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
*/
package com.easyagents.flow.core.knowledge;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.node.KnowledgeNode;
import java.util.Collections;
import java.util.ArrayList;
import java.util.List;
/**
* 工作流知识库节点的完整检索请求。
*/
public class KnowledgeSearchRequest {
private final List<Object> knowledgeIds;
private final String keyword;
private final int limit;
private final String retrievalMode;
private final KnowledgeNode knowledgeNode;
private final Chain chain;
public KnowledgeSearchRequest(
List<Object> knowledgeIds,
String keyword,
int limit,
String retrievalMode,
KnowledgeNode knowledgeNode,
Chain chain) {
this.knowledgeIds = knowledgeIds == null
? Collections.emptyList()
: Collections.unmodifiableList(new ArrayList<>(knowledgeIds));
this.keyword = keyword;
this.limit = limit;
this.retrievalMode = retrievalMode;
this.knowledgeNode = knowledgeNode;
this.chain = chain;
}
public List<Object> getKnowledgeIds() {
return knowledgeIds;
}
public String getKeyword() {
return keyword;
}
public int getLimit() {
return limit;
}
public String getRetrievalMode() {
return retrievalMode;
}
public KnowledgeNode getKnowledgeNode() {
return knowledgeNode;
}
public Chain getChain() {
return chain;
}
}

View File

@@ -15,21 +15,55 @@
*/
package com.easyagents.flow.core.node;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainSuspendException;
import com.easyagents.flow.core.chain.DataType;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.ParameterOption;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.ChainStateField;
import com.easyagents.flow.core.chain.runtime.Trigger;
import com.easyagents.flow.core.chain.runtime.TriggerContext;
import com.easyagents.flow.core.chain.runtime.TriggerType;
import java.util.*;
import java.util.ArrayList;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
public class ConfirmNode extends BaseNode {
private static final long serialVersionUID = 1L;
public static final String DEFAULT_OUTPUT_NAME = "selection";
public static final int MAX_OPTIONS = 100;
public static final int MAX_MESSAGE_LENGTH = 2000;
public static final int MAX_OPTION_LENGTH = 200;
public static final Set<String> SUPPORTED_CONFIGURATION_KEYS = Set.of(
"condition",
"description",
"expand",
"joinMode",
"loopBreakCondition",
"loopEnable",
"loopIntervalMs",
"maxLoopCount",
"maxRetryCount",
"message",
"multiple",
"options",
"outputDefs",
"resetRetryCountAfterNormal",
"retryEnable",
"retryIntervalMs",
"title");
private String message;
private List<Parameter> confirms;
private boolean multiple;
private List<String> options;
public String getMessage() {
return message;
@@ -39,115 +73,152 @@ public class ConfirmNode extends BaseNode {
this.message = message;
}
public List<Parameter> getConfirms() {
return confirms;
public boolean isMultiple() {
return multiple;
}
public void setConfirms(List<Parameter> confirms) {
if (confirms != null) {
for (Parameter confirm : confirms) {
confirm.setRefType(RefType.INPUT);
confirm.setRequired(true); // 必填,才能正确通过 getParameterValuesOnly 获取参数值
confirm.setName(confirm.getName());
}
}
this.confirms = confirms;
public void setMultiple(boolean multiple) {
this.multiple = multiple;
}
public List<String> getOptions() {
return options;
}
public void setOptions(List<String> options) {
this.options = options;
}
@Override
public Map<String, Object> execute(Chain chain) {
validateConfiguration();
String outputName = resolveOutputName();
Parameter parameter = buildParameter();
List<Parameter> confirmParameters = new ArrayList<>();
addConfirmParameter(confirmParameters);
if (confirms != null) {
for (Parameter confirm : confirms) {
Parameter clone = confirm.clone();
clone.setName(confirm.getName() + "__" + getId());
clone.setRefType(RefType.INPUT);
confirmParameters.add(clone);
// 确认值只能来自经过 Chain.resumeIfSuspended 校验后创建的恢复触发器。
// 启动参数与普通节点内存中的同名值均不能绕过人工确认。
if (!isValidatedResumeTrigger(chain)) {
chain.updateStateSafely(state -> {
if (state.getMemory().remove(parameter.getName()) == null) {
return null;
}
return EnumSet.of(ChainStateField.MEMORY);
});
}
Map<String, Object> values;
try {
values = chain.getExecutionState()
.resolveParameters(this, confirmParameters);
// 移除 confirm 参数,方便在其他节点二次确认,或者在 for 循环中第二次获取
.resolveParameters(this, Collections.singletonList(parameter));
chain.updateStateSafely(state -> {
for (Parameter confirmParameter : confirmParameters) {
state.getMemory().remove(confirmParameter.getName());
if (!state.getMemory().containsKey(parameter.getName())) {
return null;
}
state.getMemory().remove(parameter.getName());
return EnumSet.of(ChainStateField.MEMORY);
});
} catch (ChainSuspendException e) {
} catch (ChainSuspendException exception) {
chain.updateStateSafely(state -> {
state.setMessage(message);
return EnumSet.of(ChainStateField.MESSAGE);
});
if (confirms != null) {
List<Parameter> newParameters = new ArrayList<>();
for (Parameter confirm : confirms) {
Parameter clone = confirm.clone();
clone.setName(confirm.getName() + "__" + getId());
clone.setRefType(RefType.REF); // 固定为 REF
newParameters.add(clone);
throw exception;
}
// 获取参数值,不会触发 ChainSuspendException 错误
Map<String, Object> parameterValues =
chain.getExecutionState().resolveParameters(
this,
newParameters,
null,
true);
return Collections.singletonMap(
outputName,
values.get(parameter.getName()));
}
// 设置 enums方便前端给用户进行选择
for (Parameter confirmParameter : confirmParameters) {
if (confirmParameter.getEnums() == null) {
Object enumsObject = parameterValues.get(confirmParameter.getName());
confirmParameter.setEnumsObject(enumsObject);
/**
* 获取并校验当前节点配置的唯一输出名称。
*
* @return 用户配置的输出名称
*/
public String resolveOutputName() {
if (outputDefs == null || outputDefs.size() != 1
|| outputDefs.get(0) == null) {
throw new IllegalArgumentException(
"用户确认节点必须配置一个输出参数");
}
Parameter output = outputDefs.get(0);
String outputName = output.getName();
if (outputName == null || outputName.trim().isEmpty()) {
throw new IllegalArgumentException(
"用户确认节点输出参数名称不能为空");
}
DataType expectedType = multiple
? DataType.Array_String
: DataType.String;
if (output.getDataType() != expectedType) {
throw new IllegalArgumentException(
"用户确认节点输出参数类型必须与选择方式一致");
}
return outputName;
}
private boolean isValidatedResumeTrigger(Chain chain) {
Trigger trigger = TriggerContext.getCurrentTrigger();
return trigger != null
&& trigger.getType() == TriggerType.RESUME
&& Objects.equals(
chain.getStateInstanceId(),
trigger.getStateInstanceId())
&& Objects.equals(getId(), trigger.getNodeId());
}
public void validateConfiguration() {
requireText(message, MAX_MESSAGE_LENGTH, "确认提示内容");
if (options == null || options.isEmpty()) {
throw new IllegalArgumentException("用户确认节点至少需要一个选项");
}
if (options.size() > MAX_OPTIONS) {
throw new IllegalArgumentException(
"用户确认节点最多支持 " + MAX_OPTIONS + " 个选项");
}
Set<String> normalizedOptions = new HashSet<>();
for (String option : options) {
String normalized = requireText(option, MAX_OPTION_LENGTH, "选项内容");
if (!normalizedOptions.add(normalized)) {
throw new IllegalArgumentException("用户确认节点选项内容重复: " + normalized);
}
}
}
throw e;
}
Map<String, Object> results = new HashMap<>(values.size());
values.forEach((key, value) -> {
int index = key.lastIndexOf("__");
if (index >= 0) {
results.put(key.substring(0, index), value);
} else {
results.put(key, value);
}
});
return results;
}
private void addConfirmParameter(List<Parameter> parameters) {
// “确认 和 取消” 的参数
private Parameter buildParameter() {
Parameter parameter = new Parameter();
parameter.setId(DEFAULT_OUTPUT_NAME);
parameter.setName(DEFAULT_OUTPUT_NAME + "__" + getId());
parameter.setDataType(multiple ? DataType.Array_String : DataType.String);
parameter.setRefType(RefType.INPUT);
parameter.setId("confirm");
parameter.setName("confirm__" + getId());
parameter.setRequired(true);
List<Object> selectionData = new ArrayList<>();
selectionData.add("yes");
selectionData.add("no");
parameter.setEnums(selectionData);
parameter.setContentType("text");
parameter.setFormType("confirm");
parameters.add(parameter);
parameter.setFormType(multiple ? "checkbox" : "radio");
parameter.setFormLabel("选择内容");
parameter.setOptions(buildRuntimeOptions());
return parameter;
}
private List<ParameterOption> buildRuntimeOptions() {
List<ParameterOption> runtimeOptions = new ArrayList<>(options.size());
for (String option : options) {
String normalized = option.trim();
runtimeOptions.add(new ParameterOption(normalized, normalized));
}
return runtimeOptions;
}
private static String requireText(
String value, int maxLength, String fieldName) {
String normalized = value == null ? "" : value.trim();
if (normalized.isEmpty()) {
throw new IllegalArgumentException(fieldName + "不能为空");
}
if (normalized.length() > maxLength) {
throw new IllegalArgumentException(
fieldName + "不能超过 " + maxLength + " 个字符");
}
return normalized;
}
}

View File

@@ -17,15 +17,15 @@ package com.easyagents.flow.core.node;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.knowledge.Knowledge;
import com.easyagents.flow.core.knowledge.KnowledgeManager;
import com.easyagents.flow.core.util.Maps;
import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest;
import com.easyagents.flow.core.util.StringUtil;
import com.easyagents.flow.core.util.TextTemplate;
import org.slf4j.Logger;
import java.util.Arrays;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
@@ -36,6 +36,7 @@ public class KnowledgeNode extends BaseNode {
private static final Logger logger = org.slf4j.LoggerFactory.getLogger(KnowledgeNode.class);
private Object knowledgeId;
private List<Object> knowledgeIds = new ArrayList<>();
private String keyword;
private String limit;
private String retrievalMode = "HYBRID";
@@ -48,6 +49,32 @@ public class KnowledgeNode extends BaseNode {
this.knowledgeId = knowledgeId;
}
/**
* 获取规范化的知识库集合,兼容历史单值字段。
*
* @return 去重后的知识库 ID
*/
public List<Object> getKnowledgeIds() {
if (knowledgeIds != null && !knowledgeIds.isEmpty()) {
return Collections.unmodifiableList(knowledgeIds);
}
return knowledgeId == null
? Collections.emptyList()
: Collections.singletonList(knowledgeId);
}
public void setKnowledgeIds(List<?> knowledgeIds) {
LinkedHashSet<Object> normalized = new LinkedHashSet<>();
if (knowledgeIds != null) {
for (Object id : knowledgeIds) {
if (id != null && StringUtil.hasText(String.valueOf(id))) {
normalized.add(id);
}
}
}
this.knowledgeIds = new ArrayList<>(normalized);
}
public String getKeyword() {
return keyword;
}
@@ -88,25 +115,44 @@ public class KnowledgeNode extends BaseNode {
if (StringUtil.hasText(realLimitString)) {
try {
realLimit = Integer.parseInt(realLimitString);
} catch (Exception e) {
logger.error(e.toString(), e);
} catch (NumberFormatException exception) {
throw new IllegalArgumentException(
"知识库节点最终返回条数必须为正整数", exception);
}
}
Knowledge knowledge = KnowledgeManager.getInstance().getKnowledge(knowledgeId);
if (knowledge == null) {
return Collections.emptyMap();
if (realLimit <= 0) {
throw new IllegalArgumentException(
"知识库节点最终返回条数必须为正整数");
}
List<Map<String, Object>> result = knowledge.search(realKeyword, realLimit, this, chain);
return Maps.of("documents", result);
List<Object> resolvedKnowledgeIds = getKnowledgeIds();
if (resolvedKnowledgeIds.isEmpty()) {
throw new IllegalArgumentException("知识库节点至少需要选择一个知识库");
}
if (resolvedKnowledgeIds.size() > 1
&& !"VECTOR".equalsIgnoreCase(retrievalMode)) {
throw new IllegalArgumentException("多知识库检索仅支持 VECTOR 模式");
}
Map<String, Object> result = KnowledgeManager.getInstance().search(
new KnowledgeSearchRequest(
resolvedKnowledgeIds,
realKeyword,
realLimit,
retrievalMode,
this,
chain));
if (result == null) {
throw new IllegalStateException("没有可用的知识库 Provider");
}
return result;
}
@Override
public String toString() {
return "KnowledgeNode{" +
"knowledgeId=" + knowledgeId +
", knowledgeIds=" + knowledgeIds +
", keyword='" + keyword + '\'' +
", limit='" + limit + '\'' +
", retrievalMode='" + retrievalMode + '\'' +

View File

@@ -19,6 +19,8 @@ import com.alibaba.fastjson.JSON;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.WorkflowErrorReason;
import com.easyagents.flow.core.chain.WorkflowExecutionException;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.llm.LlmManager;
import com.easyagents.flow.core.util.*;
@@ -96,23 +98,21 @@ public class LlmNode extends BaseNode {
chainState.resolveParameters(this);
if (StringUtil.noText(userPrompt)) {
throw new RuntimeException("Can not find user prompt");
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID, "Can not find user prompt");
}
List<Map<String, Object>> templateRootMaps =
chainState.buildTemplateRootMaps(
parameterValues);
String userPromptString = TextTemplate.of(userPrompt)
.formatToString(templateRootMaps);
String userPromptString = formatPrompt(userPrompt, templateRootMaps);
Llm llm = LlmManager.getInstance().getChatModel(this.llmId);
if (llm == null) {
throw new RuntimeException("Can not find llm: " + this.llmId);
throw new WorkflowExecutionException(WorkflowErrorReason.MODEL_NOT_FOUND, "Can not find llm: " + this.llmId);
}
String systemPromptString = TextTemplate.of(this.systemPrompt)
.formatToString(templateRootMaps);
String systemPromptString = formatPrompt(this.systemPrompt, templateRootMaps);
Llm.MessageInfo messageInfo = new Llm.MessageInfo();
messageInfo.setMessage(userPromptString);
@@ -130,7 +130,7 @@ public class LlmNode extends BaseNode {
if (!(value instanceof String)
&& !(value instanceof java.io.File)
&& !(value instanceof Map<?, ?>)) {
throw new IllegalArgumentException(
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
"Unsupported image input for parameter '" + name + "': "
+ value.getClass().getName());
}
@@ -143,7 +143,7 @@ public class LlmNode extends BaseNode {
String responseContent = llm.chat(messageInfo, chatOptions, this, chain);
if (StringUtil.noText(responseContent)) {
throw new RuntimeException("Can not get response from llm");
throw new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID, "Can not get response from llm");
} else {
responseContent = responseContent.trim();
}
@@ -154,7 +154,10 @@ public class LlmNode extends BaseNode {
try {
jsonObjectOrArray = JSON.parse(unWrapMarkdown(responseContent));
} catch (Exception e) {
throw new RuntimeException("Can not parse json: " + responseContent + " " + e.getMessage());
throw new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID, "Can not parse model JSON output", e);
}
if (jsonObjectOrArray == null) {
throw new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID, "Model JSON output is empty");
}
if (CollectionUtil.noItems(this.outputDefs)) {
@@ -173,6 +176,14 @@ public class LlmNode extends BaseNode {
}
}
private String formatPrompt(String template, List<Map<String, Object>> values) {
try {
return TextTemplate.of(template).formatToString(values);
} catch (RuntimeException failure) {
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID, "Cannot resolve prompt variables", failure);
}
}
/**
* 移除 ``` 或者 ```json 等

View File

@@ -15,11 +15,12 @@
*/
package com.easyagents.flow.core.parser.impl;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.node.ConfirmNode;
import com.easyagents.flow.core.parser.BaseNodeParser;
import java.util.ArrayList;
import java.util.List;
public class ConfirmNodeParser extends BaseNodeParser<ConfirmNode> {
@@ -28,12 +29,36 @@ public class ConfirmNodeParser extends BaseNodeParser<ConfirmNode> {
public ConfirmNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) {
ConfirmNode confirmNode = new ConfirmNode();
confirmNode.setMessage(data.getString("message"));
List<Parameter> confirms = getParameters(data, "confirms");
if (confirms != null && !confirms.isEmpty()) {
confirmNode.setConfirms(confirms);
for (String key : data.keySet()) {
if (!ConfirmNode.SUPPORTED_CONFIGURATION_KEYS.contains(key)) {
throw new IllegalArgumentException(
"用户确认节点包含无效配置字段: " + key);
}
}
Object message = data.get("message");
if (!(message instanceof String)) {
throw new IllegalArgumentException("用户确认节点提示内容必须为字符串");
}
confirmNode.setMessage((String) message);
Object multiple = data.get("multiple");
if (!(multiple instanceof Boolean)) {
throw new IllegalArgumentException("用户确认节点选择方式必须为布尔值");
}
confirmNode.setMultiple((Boolean) multiple);
Object optionsValue = data.get("options");
if (!(optionsValue instanceof JSONArray options)) {
throw new IllegalArgumentException("用户确认节点选项必须为数组");
}
List<String> confirmOptions = new ArrayList<>(options.size());
for (Object option : options) {
if (!(option instanceof String)) {
throw new IllegalArgumentException("用户确认节点选项内容必须为字符串");
}
confirmOptions.add((String) option);
}
confirmNode.setOptions(confirmOptions);
return confirmNode;
}

View File

@@ -16,15 +16,41 @@
package com.easyagents.flow.core.parser.impl;
import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONArray;
import com.easyagents.flow.core.node.KnowledgeNode;
import com.easyagents.flow.core.parser.BaseNodeParser;
import java.util.ArrayList;
public class KnowledgeNodeParser extends BaseNodeParser<KnowledgeNode> {
@Override
public KnowledgeNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) {
KnowledgeNode knowledgeNode = new KnowledgeNode();
if (data.containsKey("knowledgeIds")) {
Object rawIds = data.get("knowledgeIds");
if (!(rawIds instanceof JSONArray)) {
throw new IllegalArgumentException("knowledgeIds 必须为数组");
}
JSONArray ids = (JSONArray) rawIds;
if (ids.isEmpty()) {
throw new IllegalArgumentException("knowledgeIds 不能为空");
}
java.util.LinkedHashSet<String> normalized =
new java.util.LinkedHashSet<>();
for (Object id : ids) {
String value = id == null ? null : String.valueOf(id).trim();
if (!com.easyagents.flow.core.util.StringUtil.hasText(value)) {
throw new IllegalArgumentException("knowledgeIds 不能包含空值");
}
if (!normalized.add(value)) {
throw new IllegalArgumentException("knowledgeIds 不能包含重复值");
}
}
knowledgeNode.setKnowledgeIds(new ArrayList<>(normalized));
} else {
knowledgeNode.setKnowledgeId(data.get("knowledgeId"));
}
knowledgeNode.setLimit(data.getString("limit"));
knowledgeNode.setKeyword(data.getString("keyword"));
knowledgeNode.setRetrievalMode(data.getString("retrievalMode"));

View File

@@ -15,10 +15,15 @@
*/
package com.easyagents.flow.core.test;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.DataType;
import com.easyagents.flow.core.chain.Edge;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.event.ChainEndEvent;
import com.easyagents.flow.core.chain.repository.ChainDefinitionSnapshotRepository;
@@ -35,6 +40,7 @@ import com.easyagents.flow.core.node.EndNode;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.node.ConfirmNode;
import com.easyagents.flow.core.node.StartNode;
import com.easyagents.flow.core.parser.ChainParser;
import org.junit.Assert;
import org.junit.Test;
@@ -61,6 +67,127 @@ import java.util.concurrent.atomic.AtomicReference;
*/
public class ChainExecutorConcurrencyTest {
/**
* 验证启动变量不能伪造确认节点的恢复参数并绕过人工确认。
*/
@Test
public void shouldSuspendConfirmNodeDespitePrefilledStartVariable()
throws Exception {
ScheduledExecutorService schedulerPool =
Executors.newSingleThreadScheduledExecutor();
ExecutorService workerPool = Executors.newFixedThreadPool(2);
TriggerScheduler triggerScheduler = new TriggerScheduler(
new InMemoryTriggerStore(), schedulerPool, workerPool, 10L);
ChainDefinition definition = createConfirmDefinition();
InMemoryChainStateRepository stateRepository =
new InMemoryChainStateRepository();
ChainExecutor executor = new ChainExecutor(
ignored -> definition,
stateRepository,
new InMemoryNodeStateRepository(),
triggerScheduler);
try {
String instanceId = executor.executeAsync(
definition.getId(),
Map.of("selection__confirm", "未配置值"));
ChainState state = awaitStatus(
stateRepository, instanceId, ChainStatus.SUSPEND);
Assert.assertFalse(
state.getMemory().containsKey("selection__confirm"));
Assert.assertEquals(
"selection__confirm",
state.getSuspendForParameters().get(0).getName());
} finally {
triggerScheduler.shutdown();
}
}
/**
* 验证设计器最终契约可解析、挂起、恢复,并把用户选择按配置名称交给结束节点。
*/
@Test
public void shouldFlowConfiguredConfirmOutputToEndNode()
throws Exception {
ScheduledExecutorService schedulerPool =
Executors.newSingleThreadScheduledExecutor();
ExecutorService workerPool = Executors.newFixedThreadPool(2);
TriggerScheduler triggerScheduler = new TriggerScheduler(
new InMemoryTriggerStore(), schedulerPool, workerPool, 10L);
ChainDefinition definition = createParsedConfirmDefinition();
InMemoryChainStateRepository stateRepository =
new InMemoryChainStateRepository();
ChainExecutor executor = new ChainExecutor(
ignored -> definition,
stateRepository,
new InMemoryNodeStateRepository(),
triggerScheduler);
try {
String instanceId = executor.executeAsync(
definition.getId(), Collections.emptyMap());
ChainState suspended = awaitStatus(
stateRepository, instanceId, ChainStatus.SUSPEND);
Assert.assertEquals(
"selection__confirm",
suspended.getSuspendForParameters().get(0).getName());
Assert.assertTrue(executor.resumeAsyncIfSuspended(
instanceId,
Map.of("selection__confirm", "继续")));
ChainState completed = awaitStatus(
stateRepository, instanceId, ChainStatus.SUCCEEDED);
Assert.assertEquals("继续", completed.getExecuteResult().get("result"));
} finally {
triggerScheduler.shutdown();
}
}
/**
* 验证多选确认结果以字符串数组形式流转到结束节点。
*/
@Test
public void shouldFlowMultipleConfirmOutputToEndNode()
throws Exception {
ScheduledExecutorService schedulerPool =
Executors.newSingleThreadScheduledExecutor();
ExecutorService workerPool = Executors.newFixedThreadPool(2);
TriggerScheduler triggerScheduler = new TriggerScheduler(
new InMemoryTriggerStore(), schedulerPool, workerPool, 10L);
ChainDefinition definition = createParsedConfirmDefinition(true);
InMemoryChainStateRepository stateRepository =
new InMemoryChainStateRepository();
ChainExecutor executor = new ChainExecutor(
ignored -> definition,
stateRepository,
new InMemoryNodeStateRepository(),
triggerScheduler);
try {
String instanceId = executor.executeAsync(
definition.getId(), Collections.emptyMap());
ChainState suspended = awaitStatus(
stateRepository, instanceId, ChainStatus.SUSPEND);
Assert.assertEquals(
DataType.Array_String,
suspended.getSuspendForParameters().get(0).getDataType());
List<String> selection = List.of("继续", "停止");
Assert.assertTrue(executor.resumeAsyncIfSuspended(
instanceId,
Map.of("selection__confirm", selection)));
ChainState completed = awaitStatus(
stateRepository, instanceId, ChainStatus.SUCCEEDED);
Assert.assertEquals(
selection, completed.getExecuteResult().get("result"));
} finally {
triggerScheduler.shutdown();
}
}
/**
* 验证同步 Tool 入口遇到人工挂起会快速失败,不会无限占用调用线程。
*
@@ -563,6 +690,12 @@ public class ChainExecutorConcurrencyTest {
start.setId("start");
ConfirmNode confirm = new ConfirmNode();
confirm.setId("confirm");
confirm.setMessage("请选择是否继续");
confirm.setOptions(List.of("继续", "停止"));
confirm.setOutputDefs(Collections.singletonList(
new Parameter(
ConfirmNode.DEFAULT_OUTPUT_NAME,
DataType.String)));
EndNode end = new EndNode();
end.setId("end");
Edge first = new Edge();
@@ -581,6 +714,91 @@ public class ChainExecutorConcurrencyTest {
return definition;
}
private ChainDefinition createParsedConfirmDefinition() {
return createParsedConfirmDefinition(false);
}
private ChainDefinition createParsedConfirmDefinition(boolean multiple) {
JSONArray nodes = new JSONArray();
nodes.add(nodeJson("start", "startNode", new JSONObject()));
JSONObject confirmData = new JSONObject();
confirmData.put("message", "请选择是否继续");
confirmData.put("multiple", multiple);
confirmData.put("options", new JSONArray(List.of("继续", "停止")));
String outputType = multiple ? "Array<String>" : "String";
confirmData.put("outputDefs", new JSONArray(List.of(
parameterJson("templateType", outputType, null))));
nodes.add(nodeJson("confirm", "confirmNode", confirmData));
JSONObject endData = new JSONObject();
endData.put("outputDefs", new JSONArray(List.of(
parameterJson(
"result", outputType, "confirm.templateType"))));
nodes.add(nodeJson("end", "endNode", endData));
JSONArray edges = new JSONArray();
edges.add(edgeJson("start-to-confirm", "start", "confirm"));
edges.add(edgeJson("confirm-to-end", "confirm", "end"));
JSONObject flow = new JSONObject();
flow.put("nodes", nodes);
flow.put("edges", edges);
ChainDefinition definition = ChainParser.builder()
.withDefaultParsers(true)
.build()
.parse(flow.toJSONString());
definition.setId("confirm-output-flow-test");
return definition;
}
private JSONObject nodeJson(
String id, String type, JSONObject data) {
JSONObject node = new JSONObject();
node.put("id", id);
node.put("type", type);
node.put("data", data);
return node;
}
private JSONObject edgeJson(
String id, String source, String target) {
JSONObject edge = new JSONObject();
edge.put("id", id);
edge.put("source", source);
edge.put("target", target);
return edge;
}
private JSONObject parameterJson(
String name, String dataType, String ref) {
JSONObject parameter = new JSONObject();
parameter.put("name", name);
parameter.put("dataType", dataType);
if (ref != null) {
parameter.put("ref", ref);
parameter.put("refType", RefType.REF.toString());
}
return parameter;
}
private ChainState awaitStatus(
InMemoryChainStateRepository repository,
String instanceId,
ChainStatus expected) throws InterruptedException {
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3);
ChainState state;
do {
state = repository.load(instanceId);
if (state != null && state.getStatus() == expected) {
return state;
}
Thread.sleep(10L);
} while (System.nanoTime() < deadline);
Assert.fail("workflow did not reach status " + expected);
return state;
}
/**
* 创建用于取消传播验证的工作流。
*

View File

@@ -2,15 +2,29 @@ package com.easyagents.flow.core.test;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainResumeException;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.EventManager;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.ParameterOption;
import com.easyagents.flow.core.chain.repository.ChainStateField;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import org.junit.Assert;
import org.junit.Test;
import java.util.Arrays;
import java.util.Collections;
import java.util.EnumSet;
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.AtomicInteger;
/**
* {@link Chain} 暂停恢复状态守卫测试。
@@ -72,6 +86,167 @@ public class ChainResumeGuardTest {
.containsKey("unexpected"));
}
/**
* 验证确认选项只能按挂起时声明的字段和值恢复。
*/
@Test
public void shouldValidateDeclaredResumeOptionsBeforeStateTransition() {
InMemoryChainStateRepository stateRepository =
new InMemoryChainStateRepository();
Chain chain = createChain(stateRepository, "resume-options");
Parameter single = optionParameter(
"templateType__confirm", "会议类型", "radio");
Parameter multiple = optionParameter(
"participants__confirm", "参会人员", "checkbox");
chain.getExecutionState().setSuspendForParameters(
Arrays.asList(single, multiple));
chain.suspend();
ChainResumeException invalidOption = assertRejected(
chain,
Map.of(
"templateType__confirm", "UNKNOWN",
"participants__confirm", List.of("REVIEW")));
ChainResumeException extraField = assertRejected(
chain,
Map.of(
"templateType__confirm", "AGENDA",
"participants__confirm", List.of("REVIEW", "REVIEW")));
assertRejected(
chain,
Map.of(
"templateType__confirm", "AGENDA",
"participants__confirm", List.of("REVIEW"),
"extra", "value"));
Assert.assertFalse(invalidOption.getMessage().contains("UNKNOWN"));
Assert.assertFalse(extraField.getMessage().contains("extra"));
Assert.assertEquals(
ChainStatus.SUSPEND,
stateRepository.load("resume-options").getStatus());
Assert.assertTrue(
stateRepository.load("resume-options").getMemory().isEmpty());
boolean resumed = chain.resumeIfSuspended(Map.of(
"templateType__confirm", "AGENDA",
"participants__confirm", List.of("REVIEW", "BRIEFING")));
Assert.assertTrue(resumed);
Assert.assertEquals(
"AGENDA",
stateRepository.load("resume-options")
.getMemory()
.get("templateType__confirm"));
Assert.assertEquals(
List.of("REVIEW", "BRIEFING"),
stateRepository.load("resume-options")
.getMemory()
.get("participants__confirm"));
}
/**
* 验证并发恢复同一暂停实例时,只有一个请求可以完成状态转换。
*/
@Test
public void shouldAllowOnlyOneConcurrentResume() throws Exception {
InMemoryChainStateRepository stateRepository =
new InMemoryChainStateRepository();
Chain first = createChain(stateRepository, "resume-concurrent");
Chain second = createChain(stateRepository, "resume-concurrent");
Parameter parameter = optionParameter(
"templateType__confirm", "会议类型", "radio");
first.getExecutionState().setSuspendForParameters(
List.of(parameter));
first.suspend();
CountDownLatch ready = new CountDownLatch(2);
CountDownLatch start = new CountDownLatch(1);
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<Boolean> firstResult = executor.submit(() -> {
ready.countDown();
start.await();
return first.resumeIfSuspended(
Map.of("templateType__confirm", "AGENDA"));
});
Future<Boolean> secondResult = executor.submit(() -> {
ready.countDown();
start.await();
return second.resumeIfSuspended(
Map.of("templateType__confirm", "REVIEW"));
});
Assert.assertTrue(ready.await(5, TimeUnit.SECONDS));
start.countDown();
int resumedCount = (firstResult.get(5, TimeUnit.SECONDS) ? 1 : 0)
+ (secondResult.get(5, TimeUnit.SECONDS) ? 1 : 0);
Assert.assertEquals(1, resumedCount);
Assert.assertEquals(
ChainStatus.RUNNING,
stateRepository.load("resume-concurrent").getStatus());
Object selected = stateRepository.load("resume-concurrent")
.getMemory()
.get("templateType__confirm");
Assert.assertTrue(
"AGENDA".equals(selected) || "REVIEW".equals(selected));
} finally {
executor.shutdownNow();
}
}
/**
* 验证恢复变量、状态和暂停上下文通过一次原子更新完成。
*/
@Test
public void shouldCommitResumeTransitionInSingleStateUpdate() {
CountingChainStateRepository stateRepository =
new CountingChainStateRepository();
Chain chain = createChain(stateRepository, "resume-single-update");
chain.getExecutionState().setSuspendForParameters(List.of(
optionParameter("templateType__confirm", "会议类型", "radio")));
chain.suspend();
stateRepository.resetUpdateCount();
boolean resumed = chain.resumeIfSuspended(
Map.of("templateType__confirm", "AGENDA"));
ChainState state = stateRepository.load("resume-single-update");
Assert.assertTrue(resumed);
Assert.assertEquals(1, stateRepository.getUpdateCount());
Assert.assertEquals(ChainStatus.RUNNING, state.getStatus());
Assert.assertNull(state.getSuspendNodeIds());
Assert.assertNull(state.getSuspendForParameters());
Assert.assertEquals("AGENDA", state.getMemory().get(
"templateType__confirm"));
}
private ChainResumeException assertRejected(
Chain chain, Map<String, Object> variables) {
try {
chain.resumeIfSuspended(variables);
Assert.fail("invalid resume variables must be rejected");
return null;
} catch (ChainResumeException expected) {
Assert.assertNotNull(expected.getMessage());
return expected;
}
}
private Parameter optionParameter(
String name, String label, String formType) {
Parameter parameter = new Parameter();
parameter.setName(name);
parameter.setFormLabel(label);
parameter.setFormType(formType);
parameter.setRequired(true);
parameter.setOptions(List.of(
new ParameterOption("第一议题", "AGENDA"),
new ParameterOption("审议类", "REVIEW"),
new ParameterOption("听取类", "BRIEFING")));
return parameter;
}
/**
* 创建使用进程内状态仓储的最小工作流。
*
@@ -93,4 +268,25 @@ public class ChainResumeGuardTest {
chain.setEventManager(new EventManager());
return chain;
}
private static class CountingChainStateRepository
extends InMemoryChainStateRepository {
private final AtomicInteger updateCount = new AtomicInteger();
@Override
public boolean tryUpdate(
ChainState chainState,
EnumSet<ChainStateField> fields) {
updateCount.incrementAndGet();
return super.tryUpdate(chainState, fields);
}
private int getUpdateCount() {
return updateCount.get();
}
private void resetUpdateCount() {
updateCount.set(0);
}
}
}

View File

@@ -0,0 +1,235 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
*/
package com.easyagents.flow.core.test;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainSuspendException;
import com.easyagents.flow.core.chain.DataType;
import com.easyagents.flow.core.chain.EventManager;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import com.easyagents.flow.core.chain.runtime.Trigger;
import com.easyagents.flow.core.chain.runtime.TriggerContext;
import com.easyagents.flow.core.chain.runtime.TriggerType;
import com.easyagents.flow.core.node.ConfirmNode;
import com.easyagents.flow.core.parser.impl.ConfirmNodeParser;
import org.junit.Assert;
import org.junit.Test;
import java.util.Collections;
import java.util.List;
import java.util.Map;
/**
* 用户确认节点选择与输出契约测试。
*/
public class ConfirmNodeTest {
@Test
public void shouldBuildSingleChoiceAndReturnSelectedContent() {
ConfirmNode node = parse(false);
node.setId("confirm-1");
Chain chain = createChain();
Parameter parameter = suspend(node, chain);
Assert.assertEquals("selection__confirm-1", parameter.getName());
Assert.assertEquals("radio", parameter.getFormType());
Assert.assertEquals("选择内容", parameter.getFormLabel());
Assert.assertEquals(DataType.String, parameter.getDataType());
Assert.assertEquals("第一议题", parameter.getOptions().get(0).getLabel());
Assert.assertEquals("第一议题", parameter.getOptions().get(0).getValue());
chain.getExecutionState().getMemory().put(parameter.getName(), "审议类");
Map<String, Object> result = executeAsResume(node, chain);
Assert.assertEquals(Collections.singletonMap("selection", "审议类"), result);
Assert.assertFalse(chain.getExecutionState().getMemory()
.containsKey(parameter.getName()));
}
@Test
public void shouldBuildMultipleChoiceAndReturnSelectedContents() {
ConfirmNode node = parse(true);
node.setId("confirm-1");
Chain chain = createChain();
Parameter parameter = suspend(node, chain);
Assert.assertEquals("checkbox", parameter.getFormType());
Assert.assertEquals(DataType.Array_String, parameter.getDataType());
List<String> selected = List.of("第一议题", "听取类");
chain.getExecutionState().getMemory().put(parameter.getName(), selected);
Assert.assertEquals(selected, executeAsResume(node, chain).get("selection"));
}
@Test
public void shouldUseConfiguredOutputNameWithoutChangingSuspendParameter() {
ConfirmNode node = parse(false, "templateChoice");
node.setId("confirm-1");
Chain chain = createChain();
Parameter parameter = suspend(node, chain);
Assert.assertEquals("selection__confirm-1", parameter.getName());
chain.getExecutionState().getMemory().put(
parameter.getName(), "第一议题");
Assert.assertEquals(
Collections.singletonMap("templateChoice", "第一议题"),
executeAsResume(node, chain));
}
@Test
public void shouldIgnorePrefilledValueWithoutResumeTrigger() {
ConfirmNode node = parse(false);
node.setId("confirm-1");
Chain chain = createChain();
chain.getExecutionState().getMemory().put(
"selection__confirm-1", "未配置值");
Parameter parameter = suspend(node, chain);
Assert.assertEquals("selection__confirm-1", parameter.getName());
Assert.assertFalse(chain.getExecutionState().getMemory()
.containsKey(parameter.getName()));
}
@Test
public void shouldRejectDuplicateNormalizedOptionContents() {
ConfirmNode node = parse(false);
node.setOptions(List.of("审议类", " 审议类 "));
try {
node.validateConfiguration();
Assert.fail("duplicate option contents must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("选项内容重复"));
}
}
@Test
public void shouldRejectNonStringOptionDuringParsing() {
JSONObject data = data(false);
data.getJSONArray("options").add(1);
try {
new ConfirmNodeParser().doParse(
new JSONObject(), data, new JSONObject());
Assert.fail("non-string option must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("必须为字符串"));
}
}
@Test
public void shouldRejectStringEncodedOptionsDuringParsing() {
JSONObject data = data(false);
data.put("options", "[\"第一议题\"]");
try {
new ConfirmNodeParser().doParse(
new JSONObject(), data, new JSONObject());
Assert.fail("string encoded options must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("必须为数组"));
}
}
@Test
public void shouldRejectImplicitModeConversionDuringParsing() {
JSONObject data = data(false);
data.put("multiple", "false");
try {
new ConfirmNodeParser().doParse(
new JSONObject(), data, new JSONObject());
Assert.fail("non-boolean mode must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("必须为布尔值"));
}
}
@Test
public void shouldRejectUnknownConfigurationFieldDuringParsing() {
JSONObject data = data(false);
data.put("async", true);
try {
new ConfirmNodeParser().doParse(
new JSONObject(), data, new JSONObject());
Assert.fail("unknown confirm configuration must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("无效配置字段"));
}
}
private static ConfirmNode parse(boolean multiple) {
return parse(multiple, ConfirmNode.DEFAULT_OUTPUT_NAME);
}
private static ConfirmNode parse(boolean multiple, String outputName) {
ConfirmNode node = new ConfirmNodeParser().doParse(
new JSONObject(), data(multiple), new JSONObject());
Parameter output = new Parameter();
output.setName(outputName);
output.setDataType(multiple
? DataType.Array_String
: DataType.String);
node.setOutputDefs(Collections.singletonList(output));
node.validateConfiguration();
return node;
}
private static JSONObject data(boolean multiple) {
JSONObject data = new JSONObject();
data.put("message", "请选择会议纪要模板");
data.put("multiple", multiple);
JSONArray options = new JSONArray();
options.addAll(List.of("第一议题", "审议类", "听取类"));
data.put("options", options);
return data;
}
private static Parameter suspend(ConfirmNode node, Chain chain) {
try {
node.execute(chain);
throw new AssertionError("confirm node must suspend");
} catch (ChainSuspendException expected) {
Assert.assertEquals(1, expected.getSuspendParameters().size());
return expected.getSuspendParameters().get(0);
}
}
private static Map<String, Object> executeAsResume(
ConfirmNode node, Chain chain) {
Trigger trigger = new Trigger();
trigger.setType(TriggerType.RESUME);
trigger.setStateInstanceId(chain.getStateInstanceId());
trigger.setNodeId(node.getId());
TriggerContext.setCurrentTrigger(trigger);
try {
return node.execute(chain);
} finally {
TriggerContext.clearCurrentTrigger();
}
}
private static Chain createChain() {
ChainDefinition definition = new ChainDefinition();
definition.setId("confirm-node-test");
definition.setNodes(Collections.emptyList());
definition.setEdges(Collections.emptyList());
Chain chain = new Chain(definition, "confirm-node-instance");
chain.setChainStateRepository(new InMemoryChainStateRepository());
chain.setNodeStateRepository(new InMemoryNodeStateRepository());
chain.setEventManager(new EventManager());
return chain;
}
}

View File

@@ -0,0 +1,217 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0.
*/
package com.easyagents.flow.core.test;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainState;
import com.easyagents.flow.core.knowledge.Knowledge;
import com.easyagents.flow.core.knowledge.KnowledgeManager;
import com.easyagents.flow.core.knowledge.KnowledgeProvider;
import com.easyagents.flow.core.knowledge.KnowledgeSearchRequest;
import com.easyagents.flow.core.node.KnowledgeNode;
import com.easyagents.flow.core.parser.impl.KnowledgeNodeParser;
import org.junit.Assert;
import org.junit.Test;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
/**
* 知识库节点多来源契约测试。
*/
public class KnowledgeNodeTest {
@Test
public void shouldParseLegacyKnowledgeId() {
JSONObject data = baseData();
data.put("knowledgeId", "101");
KnowledgeNode node = parse(data);
Assert.assertEquals(List.of("101"), node.getKnowledgeIds());
}
@Test
public void shouldPreferKnowledgeIdsAndKeepOrder() {
JSONObject data = baseData();
data.put("knowledgeId", "legacy");
JSONArray ids = new JSONArray();
ids.addAll(List.of("201", "202"));
data.put("knowledgeIds", ids);
KnowledgeNode node = parse(data);
Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds());
}
@Test
public void shouldNormalizeKnowledgeIdsBeforeStoringThem() {
JSONObject data = baseData();
JSONArray ids = new JSONArray();
ids.addAll(List.of(" 201 ", "202"));
data.put("knowledgeIds", ids);
KnowledgeNode node = parse(data);
Assert.assertEquals(List.of("201", "202"), node.getKnowledgeIds());
}
@Test
public void shouldRejectInvalidKnowledgeIds() {
JSONObject data = baseData();
data.put("knowledgeIds", "[201,202]");
assertParseFailure(data, "必须为数组");
JSONArray duplicateIds = new JSONArray();
duplicateIds.addAll(List.of("201", "201"));
data.put("knowledgeIds", duplicateIds);
assertParseFailure(data, "重复值");
}
@Test
public void defaultProviderShouldKeepSingleKnowledgeCompatibility() {
KnowledgeProvider provider = id ->
(keyword, limit, node, chain) -> List.of(Map.of(
"knowledgeId", id,
"content", keyword));
KnowledgeNode node = new KnowledgeNode();
node.setKnowledgeId("301");
Map<String, Object> output = provider.search(
new KnowledgeSearchRequest(
node.getKnowledgeIds(),
"问题",
3,
"HYBRID",
node,
null));
Assert.assertNotNull(output);
Assert.assertEquals(1, ((List<?>) output.get("documents")).size());
}
@Test
public void defaultProviderShouldDeclineMultiKnowledgeRequest() {
KnowledgeProvider provider = id -> null;
KnowledgeNode node = new KnowledgeNode();
node.setKnowledgeIds(List.of("401", "402"));
Assert.assertNull(provider.search(new KnowledgeSearchRequest(
node.getKnowledgeIds(),
"问题",
3,
"VECTOR",
node,
null)));
}
@Test
public void shouldResolveVariableLimitAndDefaultBlankValueAtRuntime() {
Assert.assertEquals(7, executeAndCaptureLimit("{{start.limit}}", "7"));
Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit}}", " "));
Assert.assertEquals(10, executeAndCaptureLimit("{{start.limit ?? }}", null));
}
@Test
public void shouldRejectInvalidResolvedVariableLimitAtRuntime() {
assertRuntimeLimitFailure("abc");
assertRuntimeLimitFailure("0");
assertRuntimeLimitFailure("-2");
}
private static KnowledgeNode parse(JSONObject data) {
return new KnowledgeNodeParser().doParse(
new JSONObject(), data, new JSONObject());
}
private static JSONObject baseData() {
JSONObject data = new JSONObject();
data.put("keyword", "问题");
data.put("limit", "5");
data.put("retrievalMode", "VECTOR");
return data;
}
private static void assertParseFailure(
JSONObject data, String expectedMessage) {
try {
parse(data);
Assert.fail("invalid knowledgeIds must be rejected");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains(expectedMessage));
}
}
private static int executeAndCaptureLimit(
String limitTemplate,
String runtimeValue) {
AtomicInteger capturedLimit = new AtomicInteger(-1);
KnowledgeProvider provider = new KnowledgeProvider() {
@Override
public Knowledge getKnowledge(Object id) {
return null;
}
@Override
public Map<String, Object> search(KnowledgeSearchRequest request) {
capturedLimit.set(request.getLimit());
return Map.of("documents", List.of());
}
};
KnowledgeManager.getInstance().registerProvider(provider);
try {
KnowledgeNode node = runtimeNode(limitTemplate);
ChainState state = new ChainState();
if (runtimeValue != null) {
state.getMemory().put("start.limit", runtimeValue);
}
node.execute(new FixedStateChain(state));
return capturedLimit.get();
} finally {
KnowledgeManager.getInstance().removeProvider(provider);
}
}
private static void assertRuntimeLimitFailure(String runtimeValue) {
KnowledgeNode node = runtimeNode("{{start.limit}}");
ChainState state = new ChainState();
state.getMemory().put("start.limit", runtimeValue);
IllegalArgumentException exception = Assert.assertThrows(
IllegalArgumentException.class,
() -> node.execute(new FixedStateChain(state)));
Assert.assertTrue(exception.getMessage().contains("必须为正整数"));
}
private static KnowledgeNode runtimeNode(String limitTemplate) {
KnowledgeNode node = new KnowledgeNode();
node.setKnowledgeIds(List.of("501", "502"));
node.setKeyword("问题");
node.setLimit(limitTemplate);
node.setRetrievalMode("VECTOR");
return node;
}
private static final class FixedStateChain extends Chain {
private final ChainState state;
private FixedStateChain(ChainState state) {
super(new ChainDefinition(), "knowledge-node-limit-test");
this.state = state;
}
@Override
public ChainState getExecutionState() {
return state;
}
}
}

View File

@@ -0,0 +1,161 @@
package com.easyagents.flow.core.test;
import com.easyagents.flow.core.chain.*;
import com.easyagents.flow.core.chain.event.*;
import com.easyagents.flow.core.chain.repository.*;
import com.easyagents.flow.core.chain.runtime.*;
import com.easyagents.flow.core.node.StartNode;
import com.easyagents.flow.core.node.LlmNode;
import com.easyagents.flow.core.llm.LlmProvider;
import com.easyagents.flow.core.llm.LlmManager;
import com.easyagents.flow.core.node.EndNode;
import org.junit.Assert;
import org.junit.Test;
import java.io.*;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicInteger;
public class WorkflowFailureContractTest {
@Test
public void failedNodeMustEndBeforeSingleTerminalEvent() throws Exception {
try (Fixture fixture = new Fixture(Integer.MAX_VALUE, false)) {
fixture.run();
Assert.assertEquals(ChainStatus.FAILED, fixture.states.load(fixture.id).getStatus());
Assert.assertEquals(List.of("FAILED", "terminal"), fixture.events);
ExceptionSummary error = fixture.states.load(fixture.id).getError();
Assert.assertEquals("worker", error.getNodeId());
Assert.assertEquals("模型分析", error.getNodeName());
Assert.assertEquals(WorkflowErrorReason.MODEL_RATE_LIMITED.getCode(), error.getErrorCode());
Assert.assertEquals(NodeStatus.FAILED, fixture.nodes.load(fixture.id, "worker").getStatus());
}
}
@Test
public void retrySuccessMustClearErrorAndKeepAttemptStates() throws Exception {
try (Fixture fixture = new Fixture(1, true)) {
fixture.run();
Assert.assertEquals(ChainStatus.SUCCEEDED, fixture.states.load(fixture.id).getStatus());
Assert.assertEquals(List.of("ERROR", "SUCCEEDED", "terminal"), fixture.events);
Assert.assertNull(fixture.nodes.load(fixture.id, "worker").getError());
Assert.assertNull(fixture.states.load(fixture.id).getError());
}
}
@Test
public void retryExhaustionMustFailWithoutSuccessEvent() throws Exception {
try (Fixture fixture = new Fixture(Integer.MAX_VALUE, true)) {
fixture.run();
Assert.assertEquals(List.of("ERROR", "FAILED", "terminal"), fixture.events);
}
}
@Test
public void summaryMustPreserveTypedCauseContextAndSerialization() throws Exception {
for (WorkflowErrorReason reason : WorkflowErrorReason.values()) {
ExceptionSummary summary = new ExceptionSummary(new RuntimeException(new WorkflowExecutionException(
reason, "internal", new IOException("raw body"))), "chain", "node", "分析");
ByteArrayOutputStream bytes = new ByteArrayOutputStream();
try (ObjectOutputStream output = new ObjectOutputStream(bytes)) { output.writeObject(summary); }
try (ObjectInputStream input = new ObjectInputStream(new ByteArrayInputStream(bytes.toByteArray()))) {
ExceptionSummary restored = (ExceptionSummary) input.readObject();
Assert.assertEquals(reason.getCode(), restored.getErrorCode());
Assert.assertEquals("node", restored.getNodeId());
Assert.assertEquals("分析", restored.getNodeName());
Assert.assertEquals(IOException.class.getName(), restored.getRootCauseClass());
}
}
Assert.assertEquals("WORKFLOW_INTERNAL_ERROR", new ExceptionSummary(new IllegalStateException()).getErrorCode());
Assert.assertEquals("NODE_EXECUTION_FAILED", new ExceptionSummary(new IllegalStateException(), "c", "n", "N").getErrorCode());
}
@Test
public void singleNodeFailureMustCarryReasonAndNode() {
try (Fixture fixture = new Fixture(1, false)) {
WorkflowExecutionException failure = Assert.assertThrows(WorkflowExecutionException.class,
() -> fixture.executor.executeNode("test", "worker", Map.of()));
Assert.assertEquals(WorkflowErrorReason.MODEL_RATE_LIMITED, failure.getReason());
Assert.assertEquals("worker", failure.getNodeId());
}
}
@Test
public void inputParsingAndModelOutputParsingMustHaveDifferentReasons() {
Chain chain = new Chain(new ChainDefinition(), "input-output-test");
chain.setChainStateRepository(new InMemoryChainStateRepository());
chain.setNodeStateRepository(new InMemoryNodeStateRepository());
ChainState state = chain.initializeState();
LlmNode node = new LlmNode(); node.setId("llm"); node.setName("模型分析");
Parameter parameter = new Parameter("items"); parameter.setRefType(RefType.REF);
parameter.setRef("items"); parameter.setDataType(DataType.Array);
state.getMemory().put("items", "not-json");
WorkflowExecutionException input = Assert.assertThrows(WorkflowExecutionException.class,
() -> state.resolveParameters(node, List.of(parameter)));
Assert.assertEquals(WorkflowErrorReason.INPUT_INVALID, input.getReason());
node.setUserPrompt("test"); node.setLlmId("error-contract-model"); node.setOutType("json");
LlmProvider provider = id -> "error-contract-model".equals(id)
? (message, options, n, c) -> "not-json" : null;
LlmManager.getInstance().registerProvider(provider);
try {
WorkflowExecutionException output = Assert.assertThrows(WorkflowExecutionException.class, () -> node.execute(chain));
Assert.assertEquals(WorkflowErrorReason.NODE_OUTPUT_INVALID, output.getReason());
Assert.assertNotNull(output.getCause());
} finally { LlmManager.getInstance().removeProvider(provider); }
}
@Test
public void missingRegisteredModelMustHaveItsOwnReason() {
Chain chain = new Chain(new ChainDefinition(), "missing-model-test");
chain.setChainStateRepository(new InMemoryChainStateRepository());
chain.setNodeStateRepository(new InMemoryNodeStateRepository());
chain.initializeState();
LlmNode node = new LlmNode();
node.setId("llm"); node.setUserPrompt("test"); node.setLlmId("absent-model-" + UUID.randomUUID());
WorkflowExecutionException error = Assert.assertThrows(WorkflowExecutionException.class, () -> node.execute(chain));
Assert.assertEquals(WorkflowErrorReason.MODEL_NOT_FOUND, error.getReason());
}
private static class Fixture implements AutoCloseable {
final InMemoryChainStateRepository states = new InMemoryChainStateRepository();
final InMemoryNodeStateRepository nodes = new InMemoryNodeStateRepository();
final TriggerScheduler scheduler = new TriggerScheduler(new InMemoryTriggerStore(),
Executors.newSingleThreadScheduledExecutor(), Executors.newFixedThreadPool(3), 1000);
final List<String> events = new CopyOnWriteArrayList<>();
final CountDownLatch ended = new CountDownLatch(1);
final ChainExecutor executor;
String id;
Fixture(int failCount, boolean retry) {
ChainDefinition definition = new ChainDefinition();
definition.setId("test");
StartNode start = new StartNode(); start.setId("start"); definition.addNode(start);
Node worker = new Node() {
final AtomicInteger attempts = new AtomicInteger();
@Override public Map<String, Object> execute(Chain chain) {
if (attempts.incrementAndGet() <= failCount) {
throw new WorkflowExecutionException(WorkflowErrorReason.MODEL_RATE_LIMITED, "internal provider body");
}
return Map.of("output", "ok");
}
};
worker.setId("worker"); worker.setName("模型分析"); worker.setRetryEnable(retry);
worker.setMaxRetryCount(1); worker.setRetryIntervalMs(5); definition.addNode(worker);
EndNode end = new EndNode(); end.setId("end"); definition.addNode(end);
for (String[] pair : List.of(new String[]{"start", "worker"}, new String[]{"worker", "end"})) {
Edge edge = new Edge(); edge.setId(pair[0] + pair[1]); edge.setSource(pair[0]); edge.setTarget(pair[1]); definition.addEdge(edge);
}
executor = new ChainExecutor(ignored -> definition, states, nodes, scheduler);
executor.addEventListener((event, chain) -> {
if (event instanceof NodeEndEvent node && node.getNode().getId().equals("worker")) events.add(node.getStatus().name());
if (event instanceof ChainStatusChangeEvent status && status.getStatus().isTerminal()) events.add("terminal");
if (event instanceof ChainEndEvent) ended.countDown();
});
}
void run() throws InterruptedException {
id = executor.executeAsync("test", Map.of());
Assert.assertTrue("workflow should finish", ended.await(5, TimeUnit.SECONDS));
}
@Override public void close() { scheduler.shutdown(); }
}
}

View File

@@ -25,7 +25,6 @@
<dependency>
<groupId>io.milvus</groupId>
<artifactId>milvus-sdk-java</artifactId>
<version>2.4.1</version>
</dependency>
<dependency>
<groupId>junit</groupId>

View File

@@ -0,0 +1,617 @@
/*
* Copyright (c) 2023-2026, Easy-Agents (fuhai999@gmail.com).
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
*/
package com.easyagents.store.milvus;
import com.easyagents.core.util.StringUtil;
import io.grpc.Context;
import io.milvus.pool.MilvusClientV2Pool;
import io.milvus.pool.PoolConfig;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.client.RetryConfig;
import java.net.URI;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.time.Duration;
import java.util.Collections;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.locks.ReentrantReadWriteLock;
import java.util.function.Function;
/**
* Shared Milvus client pool and collection state.
*/
public class MilvusClientManager implements AutoCloseable {
private static final String POOL_KEY = "default";
private static final RetryConfig SINGLE_ATTEMPT_RETRY_CONFIG = RetryConfig.builder()
.maxRetryTimes(1)
.retryOnRateLimit(false)
.maxRetryTimeoutMs(0L)
.build();
private final ReentrantReadWriteLock lifecycleLock = new ReentrantReadWriteLock();
private final Set<String> initializedCollections =
Collections.synchronizedSet(new HashSet<String>());
private final Set<String> loadedCollections =
Collections.synchronizedSet(new HashSet<String>());
private final ConcurrentMap<String, CollectionLoadTicket> collectionLoads =
new ConcurrentHashMap<String, CollectionLoadTicket>();
private final Set<Context.CancellableContext> activeContexts =
ConcurrentHashMap.newKeySet();
private final ConcurrentMap<Thread, Integer> activeOperations =
new ConcurrentHashMap<Thread, Integer>();
private volatile ManagedMilvusClientV2Pool pool;
private volatile String poolFingerprint;
private volatile long poolGeneration;
private volatile boolean acceptingOperations = true;
private volatile boolean closed;
public MilvusClientManager(MilvusVectorStoreConfig config) {
PoolSettings settings = PoolSettings.from(config);
this.pool = createPool(settings);
this.poolFingerprint = fingerprint(settings);
}
private static ManagedMilvusClientV2Pool createPool(PoolSettings settings) {
ConnectConfig connectConfig = buildConnectConfig(settings);
PoolConfig poolConfig = PoolConfig.builder()
.maxTotal(settings.poolMaxTotal())
.maxTotalPerKey(settings.poolMaxTotalPerKey())
.maxIdlePerKey(settings.poolMaxIdlePerKey())
.minIdlePerKey(settings.poolMinIdlePerKey())
.blockWhenExhausted(true)
.maxBlockWaitDuration(Duration.ofMillis(settings.poolMaxWaitMillis()))
.evictionPollingInterval(Duration.ofMillis(settings.poolEvictionIntervalMillis()))
.minEvictableIdleDuration(Duration.ofMillis(settings.poolMinEvictableIdleMillis()))
.testOnBorrow(true)
.testOnReturn(false)
.build();
try {
return new ManagedMilvusClientV2Pool(poolConfig, connectConfig);
} catch (ReflectiveOperationException exception) {
throw new IllegalStateException("Unable to initialize Milvus client pool", exception);
}
}
public <T> T withClient(Function<MilvusClientV2, T> operation) {
return withClient(null, operation);
}
public <T> T withClient(
Duration maxWait,
Function<MilvusClientV2, T> operation
) {
Thread operationThread = registerActiveOperation();
lifecycleLock.readLock().lock();
try {
ManagedMilvusClientV2Pool currentPool = requireOpenPool();
MilvusClientV2 client = maxWait == null
? currentPool.getClient(POOL_KEY)
: currentPool.getClient(POOL_KEY, maxWait);
if (client == null) {
throw new IllegalStateException(
"Milvus client pool is exhausted or unavailable"
);
}
Throwable operationFailure = null;
try {
client.retryConfig(SINGLE_ATTEMPT_RETRY_CONFIG);
Context.CancellableContext operationContext =
Context.current().withCancellation();
try {
return withRequestContext(operationContext,
() -> operation.apply(client));
} catch (RuntimeException | Error exception) {
throw exception;
} catch (Exception exception) {
throw new IllegalStateException(
"Milvus client operation failed", exception);
} finally {
operationContext.cancel(null);
}
} catch (RuntimeException | Error exception) {
operationFailure = exception;
throw exception;
} finally {
RuntimeException cleanupFailure = null;
try {
releaseClient(currentPool, client);
} catch (RuntimeException exception) {
cleanupFailure = exception;
}
if (cleanupFailure != null) {
if (operationFailure == null) {
throw cleanupFailure;
}
operationFailure.addSuppressed(cleanupFailure);
}
}
} finally {
unregisterActiveOperation(operationThread);
lifecycleLock.readLock().unlock();
}
}
private Thread registerActiveOperation() {
ensureAcceptingOperations();
Thread currentThread = Thread.currentThread();
activeOperations.merge(currentThread, 1, Integer::sum);
if (!acceptingOperations) {
unregisterActiveOperation(currentThread);
ensureAcceptingOperations();
}
return currentThread;
}
private void unregisterActiveOperation(Thread operationThread) {
activeOperations.computeIfPresent(operationThread,
(thread, depth) -> depth <= 1 ? null : depth - 1);
}
private void releaseClient(
ManagedMilvusClientV2Pool currentPool,
MilvusClientV2 client
) {
RuntimeException readinessFailure = null;
boolean reusable = false;
try {
reusable = client.clientIsReady();
} catch (RuntimeException exception) {
readinessFailure = exception;
}
try {
if (reusable) {
currentPool.returnClient(POOL_KEY, client);
} else {
discardFailedClient(currentPool, client);
}
} catch (RuntimeException cleanupFailure) {
if (readinessFailure == null) {
throw cleanupFailure;
}
readinessFailure.addSuppressed(cleanupFailure);
}
if (readinessFailure != null) {
throw readinessFailure;
}
}
<T> T withRequestContext(
Context.CancellableContext context,
Callable<T> operation
) throws Exception {
ensureAcceptingOperations();
activeContexts.add(context);
if (!acceptingOperations) {
activeContexts.remove(context);
context.cancel(new CancellationException(
"Milvus client pool is unavailable"));
ensureAcceptingOperations();
}
try {
return context.call(operation);
} finally {
activeContexts.remove(context);
}
}
private void discardFailedClient(
ManagedMilvusClientV2Pool currentPool,
MilvusClientV2 client
) {
RuntimeException cleanupFailure = null;
try {
currentPool.invalidateClient(POOL_KEY, client);
} catch (RuntimeException exception) {
cleanupFailure = exception;
}
if (cleanupFailure != null) {
throw cleanupFailure;
}
}
/**
* Rebuilds the pool when connection or pool settings change.
* Active operations are cancelled before the old pool is closed.
*
* @return true when a new pool was installed
*/
public synchronized boolean reconfigureIfNeeded(MilvusVectorStoreConfig config) {
PoolSettings nextSettings = PoolSettings.from(config);
String nextFingerprint = fingerprint(nextSettings);
if (nextFingerprint.equals(poolFingerprint)) {
return false;
}
acceptingOperations = false;
cancelActiveContexts("Milvus client pool is reconfiguring");
interruptActiveOperations();
lifecycleLock.writeLock().lock();
try {
if (closed) {
throw new IllegalStateException("Milvus client pool is closed");
}
if (nextFingerprint.equals(poolFingerprint)) {
return false;
}
ManagedMilvusClientV2Pool replacement = createPool(nextSettings);
ManagedMilvusClientV2Pool previous = pool;
pool = replacement;
poolFingerprint = nextFingerprint;
poolGeneration++;
initializedCollections.clear();
loadedCollections.clear();
failCollectionLoads("Milvus client pool was reconfigured");
if (previous != null) {
previous.close();
}
return true;
} finally {
lifecycleLock.writeLock().unlock();
if (!closed) {
acceptingOperations = true;
}
}
}
boolean isCollectionInitialized(String collectionName) {
return initializedCollections.contains(collectionName);
}
Object initializedCollectionsLock() {
return initializedCollections;
}
void markCollectionInitialized(String collectionName) {
initializedCollections.add(collectionName);
}
boolean isCollectionLoaded(String collectionName) {
return loadedCollections.contains(collectionName);
}
void markCollectionLoaded(String collectionName) {
loadedCollections.add(collectionName);
}
void markCollectionUnloaded(String collectionName) {
loadedCollections.remove(collectionName);
}
CollectionLoadTicket beginCollectionLoad(String collectionName) {
ensureAcceptingOperations();
lifecycleLock.readLock().lock();
try {
requireOpenPool();
CollectionLoadTicket candidate = new CollectionLoadTicket(
collectionName,
poolGeneration,
new CompletableFuture<Void>(),
true
);
CollectionLoadTicket existing = collectionLoads.putIfAbsent(
collectionName, candidate);
return existing == null ? candidate : existing.asFollower();
} finally {
lifecycleLock.readLock().unlock();
}
}
void completeCollectionLoad(CollectionLoadTicket ticket) {
lifecycleLock.readLock().lock();
try {
requireOpenPool();
if (ticket.generation != poolGeneration) {
throw new IllegalStateException(
"Milvus client pool changed while loading collection: "
+ ticket.collectionName
);
}
loadedCollections.add(ticket.collectionName);
ticket.completion.complete(null);
} finally {
lifecycleLock.readLock().unlock();
}
}
void failCollectionLoad(
CollectionLoadTicket ticket,
Throwable failure,
boolean retryableForFollowers
) {
if (retryableForFollowers && ticket.leader) {
collectionLoads.remove(ticket.collectionName, ticket);
}
Throwable sharedFailure = retryableForFollowers
? new RetryableCollectionLoadException(failure)
: failure;
ticket.completion.completeExceptionally(sharedFailure);
}
void endCollectionLoad(CollectionLoadTicket ticket) {
if (ticket.leader) {
collectionLoads.remove(ticket.collectionName, ticket);
}
}
public int getActiveClientCount() {
lifecycleLock.readLock().lock();
try {
return requireOpenPool().getTotalActiveClientNumber();
} finally {
lifecycleLock.readLock().unlock();
}
}
public int getIdleClientCount() {
lifecycleLock.readLock().lock();
try {
return requireOpenPool().getTotalIdleClientNumber();
} finally {
lifecycleLock.readLock().unlock();
}
}
@Override
public synchronized void close() {
if (closed) {
return;
}
acceptingOperations = false;
closed = true;
cancelActiveContexts("Milvus client pool is closing");
interruptActiveOperations();
lifecycleLock.writeLock().lock();
try {
initializedCollections.clear();
loadedCollections.clear();
poolGeneration++;
failCollectionLoads("Milvus client pool was closed");
ManagedMilvusClientV2Pool currentPool = pool;
pool = null;
poolFingerprint = null;
if (currentPool != null) {
currentPool.close();
}
} finally {
lifecycleLock.writeLock().unlock();
}
}
private ManagedMilvusClientV2Pool requireOpenPool() {
ManagedMilvusClientV2Pool currentPool = pool;
if (closed || currentPool == null) {
throw new IllegalStateException("Milvus client pool is closed");
}
return currentPool;
}
boolean isClosed() {
return closed;
}
private void ensureAcceptingOperations() {
if (!acceptingOperations) {
throw new IllegalStateException(closed
? "Milvus client pool is closed"
: "Milvus client pool is reconfiguring");
}
}
private void failCollectionLoads(String message) {
IllegalStateException failure = new IllegalStateException(message);
for (CollectionLoadTicket ticket : collectionLoads.values()) {
ticket.completion.completeExceptionally(failure);
}
collectionLoads.clear();
}
private void cancelActiveContexts(String message) {
for (Context.CancellableContext context : activeContexts) {
context.cancel(new CancellationException(message));
}
}
private void interruptActiveOperations() {
Thread currentThread = Thread.currentThread();
for (Thread operationThread : activeOperations.keySet()) {
if (operationThread != currentThread) {
operationThread.interrupt();
}
}
}
private static String fingerprint(PoolSettings settings) {
String value = String.join("\u0000",
String.valueOf(settings.uri()),
String.valueOf(settings.databaseName()),
String.valueOf(settings.token()),
String.valueOf(settings.username()),
String.valueOf(settings.password()),
String.valueOf(settings.poolMaxTotal()),
String.valueOf(settings.poolMaxTotalPerKey()),
String.valueOf(settings.poolMaxIdlePerKey()),
String.valueOf(settings.poolMinIdlePerKey()),
String.valueOf(settings.poolMaxWaitMillis()),
String.valueOf(settings.poolEvictionIntervalMillis()),
String.valueOf(settings.poolMinEvictableIdleMillis())
);
try {
byte[] digest = MessageDigest.getInstance("SHA-256")
.digest(value.getBytes(StandardCharsets.UTF_8));
StringBuilder result = new StringBuilder(digest.length * 2);
for (byte item : digest) {
result.append(String.format("%02x", item & 0xff));
}
return result.toString();
} catch (NoSuchAlgorithmException exception) {
throw new IllegalStateException("SHA-256 is unavailable", exception);
}
}
static ConnectConfig buildConnectConfig(MilvusVectorStoreConfig config) {
return buildConnectConfig(PoolSettings.from(config));
}
private static ConnectConfig buildConnectConfig(PoolSettings settings) {
String uri = normalizeAndValidateUri(settings.uri());
String databaseName = StringUtil.hasText(settings.databaseName())
? settings.databaseName().trim()
: "default";
ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder()
.uri(uri)
.dbName(databaseName);
if (StringUtil.hasText(settings.token())) {
builder.token(settings.token().trim());
}
if (StringUtil.hasText(settings.username()) && StringUtil.hasText(settings.password())) {
builder.username(settings.username().trim());
builder.password(settings.password().trim());
}
return builder.build();
}
private record PoolSettings(
String uri,
String databaseName,
String token,
String username,
String password,
int poolMaxTotal,
int poolMaxTotalPerKey,
int poolMaxIdlePerKey,
int poolMinIdlePerKey,
long poolMaxWaitMillis,
long poolEvictionIntervalMillis,
long poolMinEvictableIdleMillis
) {
private static PoolSettings from(MilvusVectorStoreConfig config) {
return new PoolSettings(
config.getUri(),
config.getDatabaseName(),
config.getToken(),
config.getUsername(),
config.getPassword(),
config.getPoolMaxTotal(),
config.getPoolMaxTotalPerKey(),
config.getPoolMaxIdlePerKey(),
config.getPoolMinIdlePerKey(),
config.getPoolMaxWaitMillis(),
config.getPoolEvictionIntervalMillis(),
config.getPoolMinEvictableIdleMillis()
);
}
}
private static final class ManagedMilvusClientV2Pool extends MilvusClientV2Pool {
private ManagedMilvusClientV2Pool(
PoolConfig poolConfig,
ConnectConfig connectConfig
) throws ClassNotFoundException, NoSuchMethodException {
super(poolConfig, connectConfig);
}
private void invalidateClient(String key, MilvusClientV2 client) {
try {
clientPool.invalidateObject(key, client);
} catch (Exception exception) {
throw new IllegalStateException("Unable to invalidate Milvus client", exception);
}
}
private MilvusClientV2 getClient(String key, Duration maxWait) {
if (maxWait == null || maxWait.isZero() || maxWait.isNegative()) {
throw new IllegalArgumentException("maxWait must be greater than zero");
}
try {
long waitMillis = Math.max(1L, maxWait.toMillis());
return clientPool.borrowObject(key, waitMillis);
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(
"Interrupted while waiting for a Milvus client", exception);
} catch (Exception exception) {
throw new IllegalStateException(
"Unable to borrow a Milvus client", exception);
}
}
}
static final class CollectionLoadTicket {
private final String collectionName;
private final long generation;
private final CompletableFuture<Void> completion;
private final boolean leader;
private CollectionLoadTicket(
String collectionName,
long generation,
CompletableFuture<Void> completion,
boolean leader
) {
this.collectionName = collectionName;
this.generation = generation;
this.completion = completion;
this.leader = leader;
}
boolean isLeader() {
return leader;
}
CompletableFuture<Void> completion() {
return completion;
}
private CollectionLoadTicket asFollower() {
return new CollectionLoadTicket(
collectionName, generation, completion, false);
}
}
static final class RetryableCollectionLoadException
extends RuntimeException {
private RetryableCollectionLoadException(Throwable cause) {
super("The collection load leader exhausted its local budget", cause);
}
}
static String normalizeAndValidateUri(String uri) {
if (StringUtil.noText(uri)) {
throw new IllegalArgumentException(
"Milvus uri is required. Example: http://127.0.0.1:19530"
);
}
String normalized = uri.trim();
if (!normalized.contains("://")) {
normalized = "http://" + normalized;
}
try {
URI parsed = URI.create(normalized);
if (StringUtil.noText(parsed.getHost()) || parsed.getPort() <= 0) {
throw new IllegalArgumentException("Invalid Milvus uri: " + uri);
}
} catch (IllegalArgumentException exception) {
throw new IllegalArgumentException(
"Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530",
exception
);
}
return normalized;
}
}

View File

@@ -15,34 +15,45 @@
*/
package com.easyagents.store.milvus;
import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.easyagents.core.document.Document;
import com.easyagents.core.store.DocumentStore;
import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult;
import com.easyagents.core.store.StoreTimeoutException;
import com.easyagents.core.util.CollectionUtil;
import com.easyagents.core.util.Maps;
import com.easyagents.core.util.StringUtil;
import io.milvus.v2.client.ConnectConfig;
import io.grpc.Context;
import io.grpc.Status;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.exception.MilvusClientException;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.vector.request.*;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import java.net.URI;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.CancellationException;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Function;
/**
* Milvus vector store based on Milvus Java SDK v2.
@@ -52,63 +63,74 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
private static final Logger LOG = LoggerFactory.getLogger(MilvusVectorStore.class);
private static final long LOAD_TIMEOUT_MS = 30_000L;
private static final long LOAD_POLL_INTERVAL_MS = 200L;
private static final long DEADLINE_SAFETY_MARGIN_MS = 200L;
private static final ScheduledExecutorService DEADLINE_SCHEDULER =
Executors.newSingleThreadScheduledExecutor(new ThreadFactory() {
@Override
public Thread newThread(Runnable runnable) {
Thread thread = new Thread(runnable, "milvus-deadline");
thread.setDaemon(true);
return thread;
}
});
private static final String FIELD_ID = "id";
private static final String FIELD_CONTENT = "content";
private static final String FIELD_METADATA = "metadata";
private static final String FIELD_VECTOR = "vector";
private final MilvusClientV2 client;
private static final Gson GSON = new Gson();
private final MilvusClientManager clientManager;
private final MilvusVectorStoreConfig config;
private final String defaultCollectionName;
private final Set<String> initializedCollections = Collections.synchronizedSet(new HashSet<String>());
private final Set<String> loadedCollections = Collections.synchronizedSet(new HashSet<String>());
private final boolean ownsClientManager;
private volatile MilvusClientV2 compatibilityClient;
private volatile boolean closed;
public MilvusVectorStore(MilvusVectorStoreConfig config) {
this(config, createOwnedClientManager(config), true);
}
public MilvusVectorStore(
MilvusVectorStoreConfig config,
MilvusClientManager clientManager
) {
this(config, clientManager, false);
}
private MilvusVectorStore(
MilvusVectorStoreConfig config,
MilvusClientManager clientManager,
boolean ownsClientManager
) {
validateConfig(config);
this.config = config;
this.defaultCollectionName = config.getDefaultCollectionName();
String uri = normalizeAndValidateUri(config.getUri());
String dbName = StringUtil.hasText(config.getDatabaseName()) ? config.getDatabaseName().trim() : "default";
ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder()
.uri(uri)
.dbName(dbName);
if (StringUtil.hasText(config.getToken())) {
builder.token(config.getToken().trim());
this.clientManager = Objects.requireNonNull(clientManager, "clientManager");
this.ownsClientManager = ownsClientManager;
}
if (StringUtil.hasText(config.getUsername()) && StringUtil.hasText(config.getPassword())) {
builder.username(config.getUsername().trim());
builder.password(config.getPassword().trim());
private static MilvusClientManager createOwnedClientManager(
MilvusVectorStoreConfig config
) {
validateConfig(config);
return new MilvusClientManager(config);
}
ConnectConfig connectConfig = builder.build();
this.client = new MilvusClientV2(connectConfig);
private static void validateConfig(MilvusVectorStoreConfig config) {
Objects.requireNonNull(config, "config");
if (config.getSearchTimeoutMillis() <= DEADLINE_SAFETY_MARGIN_MS) {
throw new IllegalArgumentException(
"Milvus searchTimeoutMillis must be greater than "
+ DEADLINE_SAFETY_MARGIN_MS
);
}
private String normalizeAndValidateUri(String uri) {
if (StringUtil.noText(uri)) {
throw new IllegalArgumentException("Milvus uri is required. Example: http://127.0.0.1:19530");
if (config.getPoolMaxWaitMillis() <= 0L) {
throw new IllegalArgumentException(
"Milvus poolMaxWaitMillis must be greater than zero"
);
}
String normalized = uri.trim();
if (!normalized.contains("://")) {
normalized = "http://" + normalized;
}
URI parsed;
try {
parsed = URI.create(normalized);
} catch (Exception e) {
throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530", e);
}
if (StringUtil.noText(parsed.getHost()) || parsed.getPort() <= 0) {
throw new IllegalArgumentException("Invalid Milvus uri: " + uri + ". Example: http://127.0.0.1:19530");
}
return normalized;
}
@Override
@@ -121,22 +143,25 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
}
int dimension = getDimension(documents);
ensureCollectionExists(collectionName, dimension);
try {
int dimension = getDimension(documents);
clientManager.withClient(client -> {
ensureCollectionExists(client, collectionName, dimension);
InsertReq.InsertReqBuilder<?, ?> builder = InsertReq.builder();
if (StringUtil.hasText(options.getPartitionName())) {
builder.partitionName(options.getPartitionName());
}
InsertReq insertReq = builder
client.insert(builder
.collectionName(collectionName)
.data(toMilvusDocuments(documents))
.build();
client.insert(insertReq);
.build());
return null;
});
return StoreResult.successWithIds(documents);
} catch (MilvusClientException e) {
return StoreResult.fail();
} catch (RuntimeException e) {
LOG.error("Milvus insert failed. collection={}, message={}",
collectionName, e.getMessage(), e);
return StoreResult.fail(e.getMessage());
}
}
@@ -159,7 +184,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
.collectionName(collectionName)
.ids(MilvusPrimaryKeySupport.normalize(ids))
.build();
clientManager.withClient(client -> {
client.delete(deleteReq);
return null;
});
return StoreResult.success();
} catch (Exception e) {
LOG.error("Milvus delete failed. collection={}, message={}",
@@ -178,19 +206,22 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
}
int dimension = getDimension(documents);
ensureCollectionExists(collectionName, dimension);
try {
UpsertReq upsertReq = UpsertReq.builder()
int dimension = getDimension(documents);
clientManager.withClient(client -> {
ensureCollectionExists(client, collectionName, dimension);
client.upsert(UpsertReq.builder()
.collectionName(collectionName)
.partitionName(options.getPartitionName())
.data(toMilvusDocuments(documents))
.build();
client.upsert(upsertReq);
.build());
return null;
});
return StoreResult.successWithIds(documents);
} catch (Exception e) {
return StoreResult.fail();
LOG.error("Milvus upsert failed. collection={}, message={}",
collectionName, e.getMessage(), e);
return StoreResult.fail(e.getMessage());
}
}
@@ -200,56 +231,104 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
if (StringUtil.noText(collectionName)) {
throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative collectionName.");
}
ensureCollectionLoaded(collectionName);
long timeoutMillis = resolveSearchTimeoutMillis(options);
long rpcBudgetMillis = timeoutMillis - DEADLINE_SAFETY_MARGIN_MS;
long deadlineNanos = deadlineAfterMillis(rpcBudgetMillis);
Context.CancellableContext context = Context.current().withDeadlineAfter(
rpcBudgetMillis, TimeUnit.MILLISECONDS, DEADLINE_SCHEDULER);
try {
return clientManager.withRequestContext(context, () ->
searchWithinDeadline(
wrapper, options, collectionName, deadlineNanos));
} catch (RuntimeException exception) {
if (!(exception instanceof StoreTimeoutException)
&& deadlineExpired(deadlineNanos, exception)) {
throw timeoutException(collectionName, exception);
}
throw exception;
} catch (Exception exception) {
throw new IllegalStateException("Milvus search failed", exception);
} finally {
context.cancel(null);
}
}
private List<Document> searchWithinDeadline(
SearchWrapper wrapper,
StoreOptions options,
String collectionName,
long deadlineNanos
) {
String operation = wrapper.getVector() == null
|| wrapper.getVector().length == 0
? "query"
: "search";
ensureCollectionLoaded(collectionName, deadlineNanos);
try {
return searchOnce(wrapper, options, collectionName, deadlineNanos);
} catch (RuntimeException exception) {
if (!isCollectionNotLoaded(exception)) {
throw propagateSearchFailure(
operation, collectionName, exception);
}
clientManager.markCollectionUnloaded(collectionName);
try {
ensureCollectionLoaded(collectionName, deadlineNanos);
return searchOnce(
wrapper, options, collectionName, deadlineNanos);
} catch (RuntimeException retryException) {
retryException.addSuppressed(exception);
throw propagateSearchFailure(
operation, collectionName, retryException);
}
}
}
private List<Document> searchOnce(
SearchWrapper wrapper,
StoreOptions options,
String collectionName,
long deadlineNanos
) {
if (wrapper.getVector() == null || wrapper.getVector().length == 0) {
return queryByCondition(wrapper, options, collectionName);
return queryByCondition(
wrapper, options, collectionName, deadlineNanos);
}
return searchByVector(wrapper, options, collectionName);
return searchByVector(wrapper, options, collectionName, deadlineNanos);
}
private List<Document> searchByVector(SearchWrapper wrapper, StoreOptions options, String collectionName) {
private List<Document> searchByVector(
SearchWrapper wrapper,
StoreOptions options,
String collectionName,
long deadlineNanos
) {
SearchReq searchReq = buildSearchReq(wrapper, options, collectionName);
try {
SearchResp resp = client.search(searchReq);
SearchResp resp = withClientBeforeDeadline(
deadlineNanos, client -> client.search(searchReq));
return parseSearchResults(resp, wrapper.getMinScore());
} catch (Exception e) {
if (isCollectionNotLoaded(e)) {
loadedCollections.remove(collectionName);
try {
ensureCollectionLoaded(collectionName);
SearchResp retryResp = client.search(searchReq);
return parseSearchResults(retryResp, wrapper.getMinScore());
} catch (Exception retryException) {
LOG.warn("Milvus search retry failed after load. collection={}, message={}", collectionName, retryException.getMessage());
return Collections.emptyList();
}
}
LOG.warn("Milvus search failed. collection={}, message={}", collectionName, e.getMessage());
return Collections.emptyList();
}
}
private List<Document> queryByCondition(SearchWrapper wrapper, StoreOptions options, String collectionName) {
private List<Document> queryByCondition(
SearchWrapper wrapper,
StoreOptions options,
String collectionName,
long deadlineNanos
) {
QueryReq queryReq = buildQueryReq(wrapper, options, collectionName);
try {
QueryResp resp = client.query(queryReq);
QueryResp resp = withClientBeforeDeadline(
deadlineNanos, client -> client.query(queryReq));
return parseQueryResults(resp);
} catch (Exception e) {
if (isCollectionNotLoaded(e)) {
loadedCollections.remove(collectionName);
try {
ensureCollectionLoaded(collectionName);
QueryResp retryResp = client.query(queryReq);
return parseQueryResults(retryResp);
} catch (Exception retryException) {
LOG.warn("Milvus query retry failed after load. collection={}, message={}", collectionName, retryException.getMessage());
return Collections.emptyList();
}
}
LOG.warn("Milvus query failed. collection={}, message={}", collectionName, e.getMessage());
return Collections.emptyList();
}
private RuntimeException propagateSearchFailure(
String operation,
String collectionName,
RuntimeException exception
) {
LOG.error("Milvus {} failed. collection={}, message={}",
operation, collectionName, exception.getMessage(), exception);
return exception;
}
private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) {
@@ -259,7 +338,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
.outputFields(getOutputFields(wrapper))
.topK(wrapper.getMaxResults())
.annsField(FIELD_VECTOR)
.data(Collections.singletonList(toFloatList(wrapper.getVector())))
.data(Collections.singletonList(new FloatVec(wrapper.getVector())))
.searchParams(Maps.of("ef", 64));
if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) {
@@ -305,11 +384,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
continue;
}
document.setId(result.getId());
Float distance = result.getDistance();
if (distance != null) {
double score = (distance + 1.0d) / 2.0d;
document.setScore(score);
}
document.setScore(normalizeScore(result.getScore()));
if (minScore == null || document.getScore() == null || document.getScore() >= minScore) {
documents.add(document);
}
@@ -318,6 +393,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
return documents;
}
static Double normalizeScore(Float rawScore) {
return rawScore == null ? null : (rawScore + 1.0d) / 2.0d;
}
private List<Document> parseQueryResults(QueryResp resp) {
List<QueryResp.QueryResult> results = resp.getQueryResults();
if (CollectionUtil.noItems(results)) {
@@ -360,22 +439,28 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
document.addMetadata(metadata);
} else if (metadataObj != null) {
@SuppressWarnings("unchecked")
Map<String, Object> metadata = JSON.parseObject(JSON.toJSONString(metadataObj), Map.class);
Map<String, Object> metadata = GSON.fromJson(
GSON.toJsonTree(metadataObj),
Map.class
);
document.addMetadata(metadata);
}
return document;
}
private List<JSONObject> toMilvusDocuments(List<Document> documents) {
List<JSONObject> rows = new ArrayList<JSONObject>(documents.size());
List<JsonObject> toMilvusDocuments(List<Document> documents) {
List<JsonObject> rows = new ArrayList<JsonObject>(documents.size());
for (Document doc : documents) {
JSONObject row = new JSONObject();
row.put(FIELD_ID, String.valueOf(doc.getId()));
row.put(FIELD_CONTENT, doc.getContent());
row.put(FIELD_VECTOR, toFloatList(doc.getVector()));
JsonObject row = new JsonObject();
row.addProperty(FIELD_ID, String.valueOf(doc.getId()));
row.addProperty(FIELD_CONTENT, doc.getContent());
row.add(FIELD_VECTOR, GSON.toJsonTree(toFloatList(doc.getVector())));
Map<String, Object> metadatas = doc.getMetadataMap();
row.put(FIELD_METADATA, metadatas == null ? new JSONObject() : new JSONObject(metadatas));
row.add(
FIELD_METADATA,
metadatas == null ? new JsonObject() : GSON.toJsonTree(metadatas)
);
rows.add(row);
}
return rows;
@@ -413,67 +498,175 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
throw new IllegalStateException("Unable to determine vector dimension for Milvus collection.");
}
private void ensureCollectionExists(String collectionName, int dimension) {
if (initializedCollections.contains(collectionName)) {
private void ensureCollectionExists(
MilvusClientV2 client,
String collectionName,
int dimension
) {
if (clientManager.isCollectionInitialized(collectionName)) {
return;
}
synchronized (initializedCollections) {
if (initializedCollections.contains(collectionName)) {
synchronized (clientManager.initializedCollectionsLock()) {
if (clientManager.isCollectionInitialized(collectionName)) {
return;
}
Boolean exists = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build());
if (Boolean.TRUE.equals(exists)) {
initializedCollections.add(collectionName);
clientManager.markCollectionInitialized(collectionName);
return;
}
if (!config.isAutoCreateCollection()) {
throw new IllegalStateException("Milvus collection not found and autoCreateCollection is disabled: " + collectionName);
}
createCollection(collectionName, dimension);
initializedCollections.add(collectionName);
createCollection(client, collectionName, dimension);
clientManager.markCollectionInitialized(collectionName);
}
}
private void ensureCollectionLoaded(String collectionName) {
if (loadedCollections.contains(collectionName)) {
private void ensureCollectionLoaded(
String collectionName,
long deadlineNanos
) {
while (!clientManager.isCollectionLoaded(collectionName)) {
MilvusClientManager.CollectionLoadTicket ticket =
clientManager.beginCollectionLoad(collectionName);
if (!ticket.isLeader()) {
if (awaitCollectionLoad(
ticket, collectionName, deadlineNanos)) {
return;
}
synchronized (loadedCollections) {
if (loadedCollections.contains(collectionName)) {
return;
continue;
}
boolean loaded = false;
try {
loaded = Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()));
} catch (Exception e) {
LOG.warn("Milvus getLoadState failed. collection={}, message={}", collectionName, e.getMessage());
if (clientManager.isCollectionLoaded(collectionName)) {
clientManager.completeCollectionLoad(ticket);
return;
}
withClientBeforeDeadline(deadlineNanos, client -> {
boolean loaded = Boolean.TRUE.equals(client.getLoadState(
GetLoadStateReq.builder()
.collectionName(collectionName)
.build()
));
if (!loaded) {
client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build());
waitForCollectionLoaded(collectionName);
client.loadCollection(LoadCollectionReq.builder()
.collectionName(collectionName)
.async(false)
.build());
waitForCollectionLoaded(
client, collectionName, deadlineNanos);
}
return null;
});
clientManager.completeCollectionLoad(ticket);
return;
} catch (RuntimeException | Error failure) {
clientManager.failCollectionLoad(
ticket, failure, isLeaderLocalAbort(failure));
throw failure;
} finally {
clientManager.endCollectionLoad(ticket);
}
loadedCollections.add(collectionName);
}
}
private void waitForCollectionLoaded(String collectionName) {
long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS;
while (System.currentTimeMillis() < deadline) {
private boolean isLeaderLocalAbort(Throwable failure) {
if (clientManager.isClosed()) {
return false;
}
if (failure instanceof StoreTimeoutException
|| Thread.currentThread().isInterrupted()
|| Context.current().isCancelled()) {
return true;
}
Throwable current = failure;
while (current != null) {
if (current instanceof InterruptedException) {
return true;
}
current = current.getCause();
}
return false;
}
private boolean awaitCollectionLoad(
MilvusClientManager.CollectionLoadTicket ticket,
String collectionName,
long deadlineNanos
) {
Context currentContext = Context.current();
CompletableFuture<Void> cancelled = new CompletableFuture<Void>();
Context.CancellationListener cancellationListener = context ->
cancelled.completeExceptionally(new CancellationException(
"Milvus search was cancelled"));
currentContext.addListener(cancellationListener, Runnable::run);
try {
CompletableFuture.anyOf(ticket.completion(), cancelled).get(
remainingNanos(deadlineNanos, collectionName),
TimeUnit.NANOSECONDS
);
return true;
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(
"Interrupted while loading Milvus collection: "
+ collectionName,
exception
);
} catch (TimeoutException exception) {
throw timeoutException(collectionName, exception);
} catch (ExecutionException exception) {
Throwable cause = exception.getCause();
if (cause instanceof MilvusClientManager
.RetryableCollectionLoadException) {
remainingNanos(deadlineNanos, collectionName);
return false;
}
if (cause instanceof RuntimeException runtimeException) {
throw runtimeException;
}
if (cause instanceof Error error) {
throw error;
}
throw new IllegalStateException(
"Unable to load Milvus collection: " + collectionName,
cause
);
} catch (CancellationException exception) {
throw new IllegalStateException(
"Milvus collection load was cancelled: " + collectionName,
exception
);
} finally {
currentContext.removeListener(cancellationListener);
}
}
private void waitForCollectionLoaded(
MilvusClientV2 client,
String collectionName,
long deadlineNanos
) {
while (true) {
long remainingNanos = remainingNanos(
deadlineNanos, collectionName);
if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) {
return;
}
try {
Thread.sleep(LOAD_POLL_INTERVAL_MS);
long sleepMillis = Math.min(
LOAD_POLL_INTERVAL_MS,
Math.max(1L, TimeUnit.NANOSECONDS.toMillis(remainingNanos))
);
Thread.sleep(sleepMillis);
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new IllegalStateException("Interrupted while loading Milvus collection: " + collectionName, e);
}
}
throw new IllegalStateException("Timeout waiting for Milvus collection loaded: " + collectionName);
}
private boolean isCollectionNotLoaded(Exception e) {
private boolean isCollectionNotLoaded(Throwable e) {
Throwable current = e;
while (current != null) {
String message = current.getMessage();
@@ -485,7 +678,91 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
return false;
}
private void createCollection(String collectionName, int dimension) {
private <T> T withClientBeforeDeadline(
long deadlineNanos,
Function<MilvusClientV2, T> operation
) {
long remainingNanos = remainingNanos(deadlineNanos, null);
long poolWaitNanos = TimeUnit.MILLISECONDS.toNanos(
config.getPoolMaxWaitMillis());
try {
return clientManager.withClient(
Duration.ofNanos(Math.min(remainingNanos, poolWaitNanos)),
operation
);
} catch (RuntimeException exception) {
if (deadlineExpired(deadlineNanos, exception)) {
throw timeoutException(null, exception);
}
throw exception;
}
}
private long resolveSearchTimeoutMillis(StoreOptions options) {
long timeoutMillis = config.getSearchTimeoutMillis();
Long requestedTimeoutMillis = options.getTimeoutMillis();
if (requestedTimeoutMillis != null) {
timeoutMillis = Math.min(timeoutMillis, requestedTimeoutMillis);
}
if (timeoutMillis <= DEADLINE_SAFETY_MARGIN_MS) {
throw new StoreTimeoutException(
"Insufficient time remaining for Milvus search"
);
}
return timeoutMillis;
}
private static long deadlineAfterMillis(long timeoutMillis) {
long now = System.nanoTime();
long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(timeoutMillis);
if (now > Long.MAX_VALUE - timeoutNanos) {
return Long.MAX_VALUE;
}
return now + timeoutNanos;
}
private static long remainingNanos(
long deadlineNanos,
String collectionName
) {
if (deadlineNanos == Long.MAX_VALUE) {
return Long.MAX_VALUE;
}
long remaining = deadlineNanos - System.nanoTime();
if (remaining <= 0L) {
throw timeoutException(collectionName, null);
}
return remaining;
}
private static boolean deadlineExpired(
long deadlineNanos,
Throwable failure
) {
if (deadlineNanos != Long.MAX_VALUE
&& System.nanoTime() >= deadlineNanos) {
return true;
}
Throwable cancellationCause = Context.current().cancellationCause();
return cancellationCause instanceof TimeoutException
|| Status.fromThrowable(failure).getCode()
== Status.Code.DEADLINE_EXCEEDED;
}
private static StoreTimeoutException timeoutException(
String collectionName,
Throwable cause
) {
String suffix = StringUtil.hasText(collectionName)
? ": " + collectionName
: "";
return new StoreTimeoutException(
"Timeout waiting for Milvus search" + suffix,
cause
);
}
private void createCollection(MilvusClientV2 client, String collectionName, int dimension) {
List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<CreateCollectionReq.FieldSchema>();
fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
.name(FIELD_ID)
@@ -531,31 +808,83 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
.indexParams(indexParams)
.build();
client.createCollection(createCollectionReq);
ensureCollectionLoaded(collectionName);
ensureCollectionLoadedForWrite(client, collectionName);
}
public MilvusClientV2 getClient() {
return client;
private void ensureCollectionLoadedForWrite(
MilvusClientV2 client,
String collectionName
) {
if (clientManager.isCollectionLoaded(collectionName)) {
return;
}
boolean loaded = Boolean.TRUE.equals(client.getLoadState(
GetLoadStateReq.builder().collectionName(collectionName).build()));
if (!loaded) {
client.loadCollection(LoadCollectionReq.builder()
.collectionName(collectionName)
.async(false)
.build());
waitForCollectionLoaded(
client,
collectionName,
deadlineAfterMillis(LOAD_TIMEOUT_MS)
);
}
clientManager.markCollectionLoaded(collectionName);
}
public boolean checkAvailable() {
try {
return client.hasCollection(HasCollectionReq.builder()
return clientManager.withClient(client -> client.hasCollection(
HasCollectionReq.builder()
.collectionName("__milvus_boot_probe__")
.build()) != null;
.build()
)) != null;
} catch (Exception e) {
LOG.warn("Milvus availability check failed. message={}", e.getMessage());
return false;
}
}
/**
* Returns a compatibility client for integrations that used the pre-pool API.
* Prefer store operations so pooled lifecycle management remains automatic.
*/
@Deprecated
public synchronized MilvusClientV2 getClient() {
if (closed) {
throw new IllegalStateException("Milvus vector store is closed");
}
if (compatibilityClient == null) {
compatibilityClient = new MilvusClientV2(
MilvusClientManager.buildConnectConfig(config)
);
}
return compatibilityClient;
}
@Override
public void close() {
MilvusClientV2 legacyClient;
synchronized (this) {
if (closed) {
return;
}
closed = true;
legacyClient = compatibilityClient;
compatibilityClient = null;
}
if (legacyClient != null) {
try {
client.close(1L);
} catch (InterruptedException e) {
legacyClient.close(1L);
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
LOG.warn("Interrupted while closing Milvus client. uri={}", config.getUri(), e);
LOG.warn("Interrupted while closing compatibility Milvus client", exception);
}
}
if (ownsClientManager) {
clientManager.close();
}
}
}

View File

@@ -30,6 +30,14 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig {
private String password;
private String defaultCollectionName;
private boolean autoCreateCollection = true;
private int poolMaxTotal = 8;
private int poolMaxTotalPerKey = 8;
private int poolMaxIdlePerKey = 4;
private int poolMinIdlePerKey = 1;
private long poolMaxWaitMillis = 3_000L;
private long poolEvictionIntervalMillis = 60_000L;
private long poolMinEvictableIdleMillis = 300_000L;
private long searchTimeoutMillis = 10_000L;
public String getUri() {
return uri;
@@ -87,6 +95,70 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig {
this.autoCreateCollection = autoCreateCollection;
}
public int getPoolMaxTotal() {
return poolMaxTotal;
}
public void setPoolMaxTotal(int poolMaxTotal) {
this.poolMaxTotal = poolMaxTotal;
}
public int getPoolMaxTotalPerKey() {
return poolMaxTotalPerKey;
}
public void setPoolMaxTotalPerKey(int poolMaxTotalPerKey) {
this.poolMaxTotalPerKey = poolMaxTotalPerKey;
}
public int getPoolMaxIdlePerKey() {
return poolMaxIdlePerKey;
}
public void setPoolMaxIdlePerKey(int poolMaxIdlePerKey) {
this.poolMaxIdlePerKey = poolMaxIdlePerKey;
}
public int getPoolMinIdlePerKey() {
return poolMinIdlePerKey;
}
public void setPoolMinIdlePerKey(int poolMinIdlePerKey) {
this.poolMinIdlePerKey = poolMinIdlePerKey;
}
public long getPoolMaxWaitMillis() {
return poolMaxWaitMillis;
}
public void setPoolMaxWaitMillis(long poolMaxWaitMillis) {
this.poolMaxWaitMillis = poolMaxWaitMillis;
}
public long getPoolEvictionIntervalMillis() {
return poolEvictionIntervalMillis;
}
public void setPoolEvictionIntervalMillis(long poolEvictionIntervalMillis) {
this.poolEvictionIntervalMillis = poolEvictionIntervalMillis;
}
public long getPoolMinEvictableIdleMillis() {
return poolMinEvictableIdleMillis;
}
public void setPoolMinEvictableIdleMillis(long poolMinEvictableIdleMillis) {
this.poolMinEvictableIdleMillis = poolMinEvictableIdleMillis;
}
public long getSearchTimeoutMillis() {
return searchTimeoutMillis;
}
public void setSearchTimeoutMillis(long searchTimeoutMillis) {
this.searchTimeoutMillis = searchTimeoutMillis;
}
@Override
public boolean checkAvailable() {
return StringUtil.hasText(this.uri);

View File

@@ -0,0 +1,92 @@
package com.easyagents.store.milvus;
import com.easyagents.core.document.Document;
import com.google.gson.JsonObject;
import org.junit.Assert;
import org.junit.Test;
import java.util.List;
import java.util.Map;
/**
* Milvus SDK 2.3.11 数据适配回归测试。
*/
public class MilvusVectorStoreCompatibilityTest {
@Test
public void shouldConvertRowsToGsonWithoutLosingMetadata() {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri("http://127.0.0.1:19530");
config.setDefaultCollectionName("test");
MilvusVectorStore store = new MilvusVectorStore(config);
try {
Document document = Document.of("正文");
document.setId("chunk-1");
document.setVector(new float[] { 0.25F, 0.75F });
document.addMetadata(Map.of("knowledgeId", "knowledge-1"));
List<JsonObject> rows = store.toMilvusDocuments(List.of(document));
Assert.assertEquals(1, rows.size());
Assert.assertEquals("chunk-1", rows.get(0).get("id").getAsString());
Assert.assertEquals("正文", rows.get(0).get("content").getAsString());
Assert.assertEquals(2, rows.get(0).getAsJsonArray("vector").size());
Assert.assertEquals(
"knowledge-1",
rows.get(0).getAsJsonObject("metadata").get("knowledgeId").getAsString()
);
} finally {
store.close();
}
}
@Test
public void shouldNormalizeUriWithoutExposingCredentialsInPoolKey() {
Assert.assertEquals(
"http://127.0.0.1:19530",
MilvusClientManager.normalizeAndValidateUri("127.0.0.1:19530")
);
}
@Test
public void shouldRebuildPoolOnlyWhenConnectionSettingsChange() {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri("http://127.0.0.1:19530");
MilvusClientManager manager = new MilvusClientManager(config);
try {
Assert.assertFalse(manager.reconfigureIfNeeded(config));
config.setPoolMaxTotal(9);
Assert.assertTrue(manager.reconfigureIfNeeded(config));
Assert.assertFalse(manager.reconfigureIfNeeded(config));
} finally {
manager.close();
}
}
@Test
public void shouldPreserveCosineScoreNormalization() {
Assert.assertEquals(Double.valueOf(1.0D), MilvusVectorStore.normalizeScore(1.0F));
Assert.assertEquals(Double.valueOf(0.5D), MilvusVectorStore.normalizeScore(0.0F));
Assert.assertEquals(Double.valueOf(0.0D), MilvusVectorStore.normalizeScore(-1.0F));
Assert.assertNull(MilvusVectorStore.normalizeScore(null));
}
@Test
public void shouldRejectCompatibilityClientAfterStoreCloses() {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri("http://127.0.0.1:19530");
MilvusVectorStore store = new MilvusVectorStore(config);
store.close();
try {
store.getClient();
Assert.fail("A closed store must not recreate a compatibility client");
} catch (IllegalStateException expected) {
Assert.assertEquals(
"Milvus vector store is closed", expected.getMessage());
}
}
}

View File

@@ -34,4 +34,24 @@ public class MilvusVectorStoreConfigTest {
config.setPassword("Milvus");
Assert.assertTrue(config.checkAvailable());
}
@Test
public void testPoolDefaultsAreBounded() {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
Assert.assertEquals(8, config.getPoolMaxTotal());
Assert.assertEquals(8, config.getPoolMaxTotalPerKey());
Assert.assertEquals(4, config.getPoolMaxIdlePerKey());
Assert.assertEquals(1, config.getPoolMinIdlePerKey());
Assert.assertEquals(3_000L, config.getPoolMaxWaitMillis());
Assert.assertEquals(300_000L, config.getPoolMinEvictableIdleMillis());
Assert.assertEquals(10_000L, config.getSearchTimeoutMillis());
}
@Test(expected = IllegalArgumentException.class)
public void testSearchTimeoutMustLeaveCleanupMargin() {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri("http://127.0.0.1:19530");
config.setSearchTimeoutMillis(200L);
new MilvusVectorStore(config);
}
}

View File

@@ -0,0 +1,935 @@
package com.easyagents.store.milvus;
import com.easyagents.core.document.Document;
import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreTimeoutException;
import io.grpc.Context;
import io.grpc.Server;
import io.grpc.netty.shaded.io.grpc.netty.NettyServerBuilder;
import io.grpc.stub.ServerCallStreamObserver;
import io.grpc.stub.StreamObserver;
import io.milvus.grpc.CheckHealthRequest;
import io.milvus.grpc.CheckHealthResponse;
import io.milvus.grpc.ConnectRequest;
import io.milvus.grpc.ConnectResponse;
import io.milvus.grpc.CollectionSchema;
import io.milvus.grpc.DataType;
import io.milvus.grpc.DescribeCollectionRequest;
import io.milvus.grpc.DescribeCollectionResponse;
import io.milvus.grpc.ErrorCode;
import io.milvus.grpc.FieldSchema;
import io.milvus.grpc.GetLoadStateRequest;
import io.milvus.grpc.GetLoadStateResponse;
import io.milvus.grpc.ListDatabasesRequest;
import io.milvus.grpc.ListDatabasesResponse;
import io.milvus.grpc.LoadCollectionRequest;
import io.milvus.grpc.LoadState;
import io.milvus.grpc.MilvusServiceGrpc;
import io.milvus.grpc.QueryRequest;
import io.milvus.grpc.QueryResults;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.service.vector.request.QueryReq;
import org.junit.Assert;
import org.junit.Test;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BooleanSupplier;
public class MilvusVectorStoreGrpcTest {
private static final long AWAIT_SECONDS = 5L;
@Test(timeout = 10_000L)
public void shouldUseOneSdkAttemptAndKeepClientReusable() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 2_000L)) {
fixture.manager.markCollectionLoaded("docs");
Assert.assertEquals(
0L,
MilvusClientManager.buildConnectConfig(fixture.config).getRpcDeadlineMs()
);
server.failQueriesWithUnavailable();
assertSearchFails(fixture.store, "docs");
Assert.assertEquals(1, server.queryCalls.get());
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
server.succeedQueries();
Assert.assertTrue(search(fixture.store, "docs").isEmpty());
Assert.assertEquals(2, server.queryCalls.get());
Assert.assertEquals(1, server.connectCalls.get());
}
}
@Test(timeout = 10_000L)
public void shouldReleaseAndReuseClientAfterContextDeadline() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 1_000L)) {
fixture.manager.markCollectionLoaded("docs");
fixture.manager.withClient(client -> client);
QueryBlock block = server.blockQueries();
ExecutorService executor = Executors.newSingleThreadExecutor();
try {
Future<List<Document>> search = executor.submit(() ->
search(fixture.store, "docs"));
block.awaitEntered();
block.awaitCancelled();
Throwable failure = futureFailure(search);
Assert.assertTrue(failure.toString(),
failure instanceof StoreTimeoutException);
Assert.assertTrue(server.querySawDeadline.get());
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
server.succeedQueries();
Assert.assertTrue(search(fixture.store, "docs").isEmpty());
Assert.assertEquals(1, server.connectCalls.get());
} finally {
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldReleaseAndReuseClientAfterThreadInterrupt() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 5_000L)) {
fixture.manager.markCollectionLoaded("docs");
fixture.manager.withClient(client -> client);
QueryBlock block = server.blockQueries();
CountDownLatch finished = new CountDownLatch(1);
AtomicReference<Throwable> failure = new AtomicReference<>();
AtomicBoolean interrupted = new AtomicBoolean();
Thread searchThread = new Thread(() -> {
try {
search(fixture.store, "docs");
} catch (Throwable exception) {
failure.set(exception);
} finally {
interrupted.set(Thread.currentThread().isInterrupted());
finished.countDown();
}
}, "milvus-interrupt-test");
searchThread.start();
block.awaitEntered();
searchThread.interrupt();
Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS));
block.awaitCancelled();
Assert.assertNotNull(failure.get());
Assert.assertTrue(interrupted.get());
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
server.succeedQueries();
Assert.assertTrue(search(fixture.store, "docs").isEmpty());
Assert.assertEquals(1, server.connectCalls.get());
}
}
@Test(timeout = 10_000L)
public void shouldKeepPoolAfterOrdinaryBusinessFailure() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 2_000L)) {
MilvusClientV2 first = fixture.manager.withClient(client -> client);
try {
fixture.manager.withClient(client -> {
throw new IllegalArgumentException("synthetic business failure");
});
Assert.fail("The business failure must be propagated");
} catch (IllegalArgumentException expected) {
Assert.assertEquals("synthetic business failure", expected.getMessage());
}
MilvusClientV2 second = fixture.manager.withClient(client -> client);
Assert.assertSame(first, second);
Assert.assertEquals(1, server.connectCalls.get());
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
}
}
@Test(timeout = 10_000L)
public void shouldAttachContextToEveryClientOperation() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 2_000L)) {
Context callerContext = Context.current();
Context operationContext = fixture.manager.withClient(client ->
Context.current());
Assert.assertNotSame(callerContext, operationContext);
Assert.assertTrue(operationContext.isCancelled());
}
}
@Test(timeout = 10_000L)
public void shouldCapPoolWaitByRemainingSearchDeadline() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 5_000L)) {
fixture.manager.markCollectionLoaded("docs");
ExecutorService executor = Executors.newSingleThreadExecutor();
CountDownLatch borrowed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
try {
Future<?> holder = executor.submit(() ->
fixture.manager.withClient(client -> {
borrowed.countDown();
try {
Assert.assertTrue(release.await(
AWAIT_SECONDS, TimeUnit.SECONDS));
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(exception);
}
return null;
}));
Assert.assertTrue(borrowed.await(
AWAIT_SECONDS, TimeUnit.SECONDS));
StoreOptions options = StoreOptions.ofCollectionName("docs");
options.setTimeoutMillis(500L);
long startedAt = System.nanoTime();
try {
search(fixture.store, options);
Assert.fail("Pool wait must respect the remaining deadline");
} catch (RuntimeException expected) {
Assert.assertTrue(expected.toString(),
expected instanceof StoreTimeoutException);
long elapsedMillis = TimeUnit.NANOSECONDS.toMillis(
System.nanoTime() - startedAt);
Assert.assertTrue("elapsedMillis=" + elapsedMillis,
elapsedMillis < 800L);
}
release.countDown();
holder.get(AWAIT_SECONDS, TimeUnit.SECONDS);
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
} finally {
release.countDown();
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldInvalidateOnlyClosedClient() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 2, 2_000L)) {
ExecutorService executor = Executors.newSingleThreadExecutor();
CountDownLatch firstBorrowed = new CountDownLatch(1);
CountDownLatch releaseFirst = new CountDownLatch(1);
try {
Future<?> holder = executor.submit(() ->
fixture.manager.withClient(client -> {
firstBorrowed.countDown();
try {
Assert.assertTrue(releaseFirst.await(
AWAIT_SECONDS, TimeUnit.SECONDS));
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(exception);
}
return null;
}));
Assert.assertTrue(firstBorrowed.await(
AWAIT_SECONDS, TimeUnit.SECONDS));
fixture.manager.withClient(client -> client);
releaseFirst.countDown();
holder.get(AWAIT_SECONDS, TimeUnit.SECONDS);
Assert.assertEquals(2, fixture.manager.getIdleClientCount());
AtomicReference<MilvusClientV2> closed = new AtomicReference<>();
try {
fixture.manager.withClient(client -> {
closed.set(client);
client.close();
throw new IllegalStateException("synthetic closed client");
});
Assert.fail("The closed-client failure must be propagated");
} catch (IllegalStateException expected) {
Assert.assertEquals(
"synthetic closed client", expected.getMessage());
}
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
MilvusClientV2 remaining = fixture.manager.withClient(
client -> client);
Assert.assertNotSame(closed.get(), remaining);
Assert.assertEquals(2, server.connectCalls.get());
} finally {
releaseFirst.countDown();
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldLoadSameCollectionOnlyOnce() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 2, 3_000L)) {
LoadGate gate = server.blockLoad("shared");
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<List<Document>> first = executor.submit(() ->
search(fixture.store, "shared"));
gate.awaitEntered();
Future<List<Document>> second = executor.submit(() ->
search(fixture.store, "shared"));
MilvusClientManager.CollectionLoadTicket follower =
fixture.manager.beginCollectionLoad("shared");
Assert.assertFalse(follower.isLeader());
awaitCondition(() -> follower.completion().getNumberOfDependents() > 0);
Assert.assertEquals(1, server.loadCalls("shared"));
Assert.assertEquals(1, fixture.manager.getActiveClientCount());
gate.release();
Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
Assert.assertEquals(1, server.loadCalls("shared"));
Assert.assertTrue(fixture.manager.isCollectionLoaded("shared"));
} finally {
gate.release();
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldRemoveLeaderTicketAfterLateCacheHit() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer()) {
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri(server.uri());
config.setDefaultCollectionName("docs");
config.setPoolMinIdlePerKey(0);
config.setSearchTimeoutMillis(2_000L);
RacingMilvusClientManager manager =
new RacingMilvusClientManager(config, "late-hit");
MilvusVectorStore store = new MilvusVectorStore(config, manager);
try {
Assert.assertTrue(search(store, "late-hit").isEmpty());
manager.markCollectionUnloaded("late-hit");
Assert.assertTrue(search(store, "late-hit").isEmpty());
Assert.assertEquals(1, server.loadCalls("late-hit"));
Assert.assertTrue(manager.isCollectionLoaded("late-hit"));
} finally {
store.close();
manager.close();
}
}
}
@Test(timeout = 10_000L)
public void shouldLoadDifferentCollectionsInParallel() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 2, 3_000L)) {
LoadGate firstGate = server.blockLoad("first");
LoadGate secondGate = server.blockLoad("second");
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<List<Document>> first = executor.submit(() ->
search(fixture.store, "first"));
Future<List<Document>> second = executor.submit(() ->
search(fixture.store, "second"));
firstGate.awaitEntered();
secondGate.awaitEntered();
Assert.assertEquals(2, server.activeLoads.get());
Assert.assertEquals(2, server.maxConcurrentLoads.get());
Assert.assertEquals(2, fixture.manager.getActiveClientCount());
firstGate.release();
secondGate.release();
Assert.assertTrue(first.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
Assert.assertTrue(second.get(AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
} finally {
firstGate.release();
secondGate.release();
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldReelectFollowerAfterLoadLeaderTimesOut() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 3_000L)) {
LoadGate gate = server.blockLoad("reelect");
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
StoreOptions shortBudget =
StoreOptions.ofCollectionName("reelect");
shortBudget.setTimeoutMillis(500L);
Future<List<Document>> first = executor.submit(() ->
search(fixture.store, shortBudget));
gate.awaitEntered();
Future<List<Document>> second = executor.submit(() ->
search(fixture.store, "reelect"));
awaitCondition(() -> server.loadCalls("reelect") == 2);
gate.release();
Throwable firstFailure = futureFailure(first);
Assert.assertTrue(firstFailure.toString(),
firstFailure instanceof StoreTimeoutException);
Assert.assertTrue(second.get(
AWAIT_SECONDS, TimeUnit.SECONDS).isEmpty());
Assert.assertEquals(2, server.loadCalls("reelect"));
Assert.assertTrue(
fixture.manager.isCollectionLoaded("reelect"));
} finally {
gate.release();
executor.shutdownNow();
}
}
}
@Test(timeout = 10_000L)
public void shouldCancelWaitingFollowerWithoutBorrowingClient() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 5_000L)) {
MilvusClientManager.CollectionLoadTicket leader =
fixture.manager.beginCollectionLoad("waiting");
CountDownLatch finished = new CountDownLatch(1);
AtomicReference<Throwable> failure = new AtomicReference<>();
Thread follower = new Thread(() -> {
try {
search(fixture.store, "waiting");
} catch (Throwable exception) {
failure.set(exception);
} finally {
finished.countDown();
}
}, "milvus-load-follower-test");
try {
follower.start();
awaitCondition(() -> leader.completion().getNumberOfDependents() > 0);
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(0, server.connectCalls.get());
follower.interrupt();
Assert.assertTrue(finished.await(AWAIT_SECONDS, TimeUnit.SECONDS));
Assert.assertNotNull(failure.get());
Assert.assertFalse(leader.completion().isDone());
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(0, server.connectCalls.get());
} finally {
follower.interrupt();
fixture.manager.failCollectionLoad(
leader, new IllegalStateException("test cleanup"), false);
fixture.manager.endCollectionLoad(leader);
}
}
}
@Test(timeout = 10_000L)
public void shouldRecoverAfterCollectionLoadFailure() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 2_000L)) {
server.failNextLoad("recoverable");
LoadGate gate = server.blockLoad("recoverable");
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
Future<List<Document>> first = executor.submit(() ->
search(fixture.store, "recoverable"));
gate.awaitEntered();
Future<List<Document>> second = executor.submit(() ->
search(fixture.store, "recoverable"));
MilvusClientManager.CollectionLoadTicket follower =
fixture.manager.beginCollectionLoad("recoverable");
Assert.assertFalse(follower.isLeader());
awaitCondition(() ->
follower.completion().getNumberOfDependents() > 0);
gate.release();
assertFutureFails(first);
assertFutureFails(second);
} finally {
gate.release();
executor.shutdownNow();
}
Assert.assertEquals(1, server.loadCalls("recoverable"));
Assert.assertFalse(fixture.manager.isCollectionLoaded("recoverable"));
Assert.assertEquals(0, fixture.manager.getActiveClientCount());
Assert.assertEquals(1, fixture.manager.getIdleClientCount());
Assert.assertTrue(search(fixture.store, "recoverable").isEmpty());
Assert.assertEquals(2, server.loadCalls("recoverable"));
Assert.assertTrue(fixture.manager.isCollectionLoaded("recoverable"));
Assert.assertEquals(1, server.connectCalls.get());
}
}
@Test(timeout = 10_000L)
public void shouldCancelActiveRpcWhenManagerCloses() throws Exception {
FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 5_000L);
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
fixture.manager.markCollectionLoaded("docs");
QueryBlock block = server.blockQueries();
Future<?> operation = executor.submit(() ->
query(fixture.manager));
block.awaitEntered();
Future<?> close = executor.submit(fixture.manager::close);
block.awaitCancelled();
close.get(AWAIT_SECONDS, TimeUnit.SECONDS);
assertFutureFails(operation);
try {
fixture.manager.withClient(client -> null);
Assert.fail("A closed manager must reject client borrows");
} catch (IllegalStateException expected) {
Assert.assertEquals("Milvus client pool is closed", expected.getMessage());
}
} finally {
executor.shutdownNow();
fixture.close();
server.close();
}
}
@Test(timeout = 10_000L)
public void shouldCancelActiveRpcWhenManagerReconfigures() throws Exception {
try (FakeMilvusServer server = new FakeMilvusServer();
Fixture fixture = new Fixture(server, 1, 5_000L)) {
ExecutorService executor = Executors.newFixedThreadPool(2);
try {
QueryBlock block = server.blockQueries();
Future<?> operation = executor.submit(() ->
query(fixture.manager));
block.awaitEntered();
fixture.config.setPoolMaxTotal(2);
Future<Boolean> reconfigure = executor.submit(() ->
fixture.manager.reconfigureIfNeeded(fixture.config));
block.awaitCancelled();
Assert.assertTrue(reconfigure.get(
AWAIT_SECONDS, TimeUnit.SECONDS));
assertFutureFails(operation);
server.succeedQueries();
query(fixture.manager);
} finally {
executor.shutdownNow();
}
}
}
private static Object query(MilvusClientManager manager) {
return manager.withClient(client -> client.query(
QueryReq.builder()
.collectionName("docs")
.filter("id == \"synthetic-id\"")
.outputFields(List.of("id"))
.build()
));
}
private static List<Document> search(MilvusVectorStore store, String collection) {
return search(store, StoreOptions.ofCollectionName(collection));
}
private static List<Document> search(
MilvusVectorStore store,
StoreOptions options
) {
SearchWrapper wrapper = new SearchWrapper();
wrapper.setWithVector(false);
wrapper.eq("id", "synthetic-id");
return store.search(wrapper, options);
}
private static void assertSearchFails(MilvusVectorStore store, String collection) {
try {
search(store, collection);
Assert.fail("The synthetic Milvus failure must be propagated");
} catch (RuntimeException expected) {
Assert.assertNotNull(expected);
}
}
private static void assertFutureFails(Future<?> future)
throws InterruptedException, TimeoutException {
futureFailure(future);
}
private static Throwable futureFailure(Future<?> future)
throws InterruptedException, TimeoutException {
try {
future.get(AWAIT_SECONDS, TimeUnit.SECONDS);
Assert.fail("The synthetic Milvus failure must be propagated");
} catch (ExecutionException expected) {
Assert.assertNotNull(expected.getCause());
return expected.getCause();
}
throw new AssertionError("Expected future to fail");
}
private static void awaitCondition(BooleanSupplier condition)
throws InterruptedException, TimeoutException {
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(AWAIT_SECONDS);
while (!condition.getAsBoolean()) {
if (System.nanoTime() >= deadline) {
throw new TimeoutException("Timed out waiting for test condition");
}
if (Thread.interrupted()) {
throw new InterruptedException();
}
Thread.onSpinWait();
}
}
private static final class Fixture implements AutoCloseable {
private final MilvusVectorStoreConfig config;
private final MilvusClientManager manager;
private final MilvusVectorStore store;
private Fixture(FakeMilvusServer server, int poolSize, long searchTimeoutMillis) {
config = new MilvusVectorStoreConfig();
config.setUri(server.uri());
config.setDefaultCollectionName("docs");
config.setPoolMaxTotal(poolSize);
config.setPoolMaxTotalPerKey(poolSize);
config.setPoolMaxIdlePerKey(poolSize);
config.setPoolMinIdlePerKey(0);
config.setPoolMaxWaitMillis(1_000L);
config.setSearchTimeoutMillis(searchTimeoutMillis);
manager = new MilvusClientManager(config);
store = new MilvusVectorStore(config, manager);
}
@Override
public void close() {
store.close();
manager.close();
}
}
private static final class RacingMilvusClientManager
extends MilvusClientManager {
private final String collectionName;
private final AtomicInteger observations = new AtomicInteger();
private RacingMilvusClientManager(
MilvusVectorStoreConfig config,
String collectionName
) {
super(config);
this.collectionName = collectionName;
}
@Override
boolean isCollectionLoaded(String requestedCollectionName) {
if (collectionName.equals(requestedCollectionName)) {
int observation = observations.getAndIncrement();
if (observation == 0) {
return false;
}
if (observation == 1) {
return true;
}
}
return super.isCollectionLoaded(requestedCollectionName);
}
}
private static final class QueryBlock {
private final CountDownLatch entered = new CountDownLatch(1);
private final CountDownLatch cancelled = new CountDownLatch(1);
private void awaitEntered() throws InterruptedException {
Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS));
}
private void awaitCancelled() throws InterruptedException {
Assert.assertTrue(cancelled.await(AWAIT_SECONDS, TimeUnit.SECONDS));
}
}
private static final class LoadGate {
private final CountDownLatch entered = new CountDownLatch(1);
private final CountDownLatch release = new CountDownLatch(1);
private void awaitEntered() throws InterruptedException {
Assert.assertTrue(entered.await(AWAIT_SECONDS, TimeUnit.SECONDS));
}
private void release() {
release.countDown();
}
}
private static final class FakeMilvusServer implements AutoCloseable {
private static final io.milvus.grpc.Status SUCCESS =
io.milvus.grpc.Status.newBuilder()
.setErrorCode(ErrorCode.Success)
.setCode(0)
.build();
private final AtomicInteger connectCalls = new AtomicInteger();
private final AtomicInteger queryCalls = new AtomicInteger();
private final AtomicInteger activeLoads = new AtomicInteger();
private final AtomicInteger maxConcurrentLoads = new AtomicInteger();
private final AtomicBoolean querySawDeadline = new AtomicBoolean();
private final ConcurrentHashMap<String, AtomicInteger> loadCalls =
new ConcurrentHashMap<>();
private final ConcurrentHashMap<String, LoadGate> loadGates =
new ConcurrentHashMap<>();
private final Set<String> loadedCollections = ConcurrentHashMap.newKeySet();
private final Set<String> failNextLoads = ConcurrentHashMap.newKeySet();
private final ExecutorService rpcExecutor = Executors.newCachedThreadPool();
private final Server server;
private volatile QueryAction queryAction = QueryAction.SUCCESS;
private volatile QueryBlock queryBlock;
private FakeMilvusServer() throws IOException {
server = NettyServerBuilder.forPort(0)
.executor(rpcExecutor)
.addService(new Service())
.build()
.start();
}
private String uri() {
return "http://127.0.0.1:" + server.getPort();
}
private void succeedQueries() {
queryAction = QueryAction.SUCCESS;
queryBlock = null;
}
private void failQueriesWithUnavailable() {
queryAction = QueryAction.UNAVAILABLE;
queryBlock = null;
}
private QueryBlock blockQueries() {
QueryBlock block = new QueryBlock();
queryBlock = block;
queryAction = QueryAction.BLOCK;
return block;
}
private LoadGate blockLoad(String collectionName) {
LoadGate gate = new LoadGate();
loadGates.put(collectionName, gate);
return gate;
}
private void failNextLoad(String collectionName) {
failNextLoads.add(collectionName);
}
private int loadCalls(String collectionName) {
AtomicInteger calls = loadCalls.get(collectionName);
return calls == null ? 0 : calls.get();
}
@Override
public void close() {
for (LoadGate gate : new ArrayList<>(loadGates.values())) {
gate.release();
}
server.shutdownNow();
try {
server.awaitTermination(AWAIT_SECONDS, TimeUnit.SECONDS);
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
} finally {
rpcExecutor.shutdownNow();
}
}
private enum QueryAction {
SUCCESS,
UNAVAILABLE,
BLOCK
}
private final class Service extends MilvusServiceGrpc.MilvusServiceImplBase {
@Override
public void connect(
ConnectRequest request,
StreamObserver<ConnectResponse> observer
) {
connectCalls.incrementAndGet();
observer.onNext(ConnectResponse.newBuilder()
.setStatus(SUCCESS)
.setIdentifier(1L)
.build());
observer.onCompleted();
}
@Override
public void listDatabases(
ListDatabasesRequest request,
StreamObserver<ListDatabasesResponse> observer
) {
observer.onNext(ListDatabasesResponse.newBuilder()
.setStatus(SUCCESS)
.addDbNames("default")
.build());
observer.onCompleted();
}
@Override
public void checkHealth(
CheckHealthRequest request,
StreamObserver<CheckHealthResponse> observer
) {
observer.onNext(CheckHealthResponse.newBuilder()
.setStatus(SUCCESS)
.setIsHealthy(true)
.build());
observer.onCompleted();
}
@Override
public void getLoadState(
GetLoadStateRequest request,
StreamObserver<GetLoadStateResponse> observer
) {
LoadState state = loadedCollections.contains(request.getCollectionName())
? LoadState.LoadStateLoaded
: LoadState.LoadStateNotLoad;
observer.onNext(GetLoadStateResponse.newBuilder()
.setStatus(SUCCESS)
.setState(state)
.build());
observer.onCompleted();
}
@Override
public void describeCollection(
DescribeCollectionRequest request,
StreamObserver<DescribeCollectionResponse> observer
) {
CollectionSchema schema = CollectionSchema.newBuilder()
.setName(request.getCollectionName())
.addFields(FieldSchema.newBuilder()
.setName("id")
.setIsPrimaryKey(true)
.setDataType(DataType.VarChar)
.build())
.build();
observer.onNext(DescribeCollectionResponse.newBuilder()
.setStatus(SUCCESS)
.setCollectionName(request.getCollectionName())
.setSchema(schema)
.build());
observer.onCompleted();
}
@Override
public void loadCollection(
LoadCollectionRequest request,
StreamObserver<io.milvus.grpc.Status> observer
) {
String collectionName = request.getCollectionName();
loadCalls.computeIfAbsent(
collectionName, ignored -> new AtomicInteger()).incrementAndGet();
LoadGate gate = loadGates.get(collectionName);
if (gate != null) {
int active = activeLoads.incrementAndGet();
maxConcurrentLoads.accumulateAndGet(active, Math::max);
gate.entered.countDown();
try {
if (!gate.release.await(AWAIT_SECONDS, TimeUnit.SECONDS)) {
observer.onError(io.grpc.Status.DEADLINE_EXCEEDED
.withDescription("test load gate timed out")
.asRuntimeException());
return;
}
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
observer.onError(io.grpc.Status.CANCELLED
.withCause(exception)
.asRuntimeException());
return;
} finally {
activeLoads.decrementAndGet();
}
}
if (failNextLoads.remove(collectionName)) {
observer.onNext(io.milvus.grpc.Status.newBuilder()
.setErrorCode(ErrorCode.UnexpectedError)
.setCode(1)
.setReason("synthetic load failure")
.build());
observer.onCompleted();
return;
}
loadedCollections.add(collectionName);
observer.onNext(SUCCESS);
observer.onCompleted();
}
@Override
public void query(
QueryRequest request,
StreamObserver<QueryResults> observer
) {
queryCalls.incrementAndGet();
querySawDeadline.compareAndSet(
false, Context.current().getDeadline() != null);
QueryAction action = queryAction;
if (action == QueryAction.UNAVAILABLE) {
observer.onError(io.grpc.Status.UNAVAILABLE
.withDescription("synthetic query failure")
.asRuntimeException());
return;
}
if (action == QueryAction.BLOCK) {
QueryBlock block = queryBlock;
@SuppressWarnings("unchecked")
ServerCallStreamObserver<QueryResults> serverObserver =
(ServerCallStreamObserver<QueryResults>) observer;
serverObserver.setOnCancelHandler(block.cancelled::countDown);
block.entered.countDown();
return;
}
observer.onNext(QueryResults.newBuilder()
.setStatus(SUCCESS)
.setCollectionName(request.getCollectionName())
.build());
observer.onCompleted();
}
}
}
}

View File

@@ -0,0 +1,190 @@
package com.easyagents.store.milvus;
import com.easyagents.core.document.Document;
import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
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.AtomicReference;
/**
* Opt-in compatibility smoke tests for a real Milvus instance.
*/
public class MilvusVectorStoreIntegrationTest {
@Test
public void shouldCrudAgainstRealMilvusAndRecoverPooledClients() throws Exception {
String uri = System.getenv("MILVUS_TEST_URI");
Assume.assumeTrue("MILVUS_TEST_URI is not configured", uri != null && !uri.isBlank());
String collectionName = "easy_agents_sdk_2311_" + UUID.randomUUID().toString().replace("-", "");
MilvusVectorStoreConfig config = new MilvusVectorStoreConfig();
config.setUri(uri);
config.setDefaultCollectionName(collectionName);
config.setPoolMaxTotal(1);
config.setPoolMaxTotalPerKey(1);
config.setPoolMaxIdlePerKey(1);
config.setPoolMinIdlePerKey(0);
config.setPoolMaxWaitMillis(250L);
MilvusClientManager manager = new MilvusClientManager(config);
MilvusVectorStore store = new MilvusVectorStore(config, manager);
StoreOptions options = StoreOptions.ofCollectionName(collectionName);
try {
Document first = document("chunk-1", "first", 1.0F, 0.0F);
Document second = document("chunk-2", "second", 0.0F, 1.0F);
Assert.assertTrue(store.store(List.of(first, second), options).isSuccess());
SearchWrapper nearest = new SearchWrapper();
nearest.setVector(new float[] { 1.0F, 0.0F });
nearest.setMaxResults(1);
List<Document> initial = store.search(nearest, options);
Assert.assertEquals(1, initial.size());
Assert.assertEquals("chunk-1", String.valueOf(initial.get(0).getId()));
Document updated = document("chunk-1", "updated", 1.0F, 0.0F);
Assert.assertTrue(store.update(List.of(updated), options).isSuccess());
Assert.assertEquals("updated", store.search(nearest, options).get(0).getContent());
Assert.assertTrue(store.delete(List.of("chunk-2"), options).isSuccess());
SearchWrapper deleted = new SearchWrapper();
deleted.setWithVector(false);
deleted.eq("id", "chunk-2");
Assert.assertTrue(store.search(deleted, options).isEmpty());
assertQueryFailureIsNotReportedAsEmpty(store, options);
assertPoolExhaustionIsBounded(manager);
assertBusinessFailurePreservesClient(manager);
Assert.assertTrue(store.checkAvailable());
} finally {
try {
manager.withClient(client -> {
client.dropCollection(DropCollectionReq.builder()
.collectionName(collectionName)
.build());
return null;
});
} finally {
store.close();
manager.close();
}
}
try {
manager.getActiveClientCount();
Assert.fail("A closed pool must reject further use");
} catch (IllegalStateException expected) {
Assert.assertEquals("Milvus client pool is closed", expected.getMessage());
}
}
private static void assertQueryFailureIsNotReportedAsEmpty(
MilvusVectorStore store,
StoreOptions options
) {
SearchWrapper invalid = new SearchWrapper();
invalid.setWithVector(false);
invalid.eq("id", "chunk-1");
StoreOptions invalidOptions = StoreOptions.ofCollectionName(
options.getCollectionName()
).partitionName("__missing_partition__");
try {
store.search(invalid, invalidOptions);
Assert.fail("A Milvus query failure must not be reported as an empty result");
} catch (RuntimeException expected) {
Assert.assertNotNull(expected.getMessage());
}
}
private static void assertBusinessFailurePreservesClient(MilvusClientManager manager)
throws InterruptedException, ExecutionException {
ExecutorService executor = Executors.newFixedThreadPool(2);
CountDownLatch borrowed = new CountDownLatch(1);
CountDownLatch waiterStarted = new CountDownLatch(1);
CountDownLatch fail = new CountDownLatch(1);
AtomicReference<Object> failedClient = new AtomicReference<>();
try {
Future<?> failing = executor.submit(() -> {
try {
manager.withClient(client -> {
failedClient.set(client);
borrowed.countDown();
try {
if (!fail.await(2, TimeUnit.SECONDS)) {
throw new IllegalStateException("Timed out waiting to fail client");
}
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(exception);
}
throw new IllegalStateException("synthetic RPC failure");
});
Assert.fail("The synthetic client failure must be propagated");
} catch (IllegalStateException expected) {
Assert.assertEquals("synthetic RPC failure", expected.getMessage());
}
});
Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS));
Future<Object> waiting = executor.submit(() -> {
waiterStarted.countDown();
return manager.withClient(client -> client);
});
Assert.assertTrue(waiterStarted.await(2, TimeUnit.SECONDS));
fail.countDown();
failing.get();
Assert.assertSame(failedClient.get(), waiting.get());
} finally {
fail.countDown();
executor.shutdownNow();
}
}
private static void assertPoolExhaustionIsBounded(MilvusClientManager manager)
throws InterruptedException, ExecutionException {
ExecutorService executor = Executors.newFixedThreadPool(2);
CountDownLatch borrowed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
try {
Future<?> holder = executor.submit(() -> manager.withClient(client -> {
borrowed.countDown();
try {
if (!release.await(2, TimeUnit.SECONDS)) {
throw new IllegalStateException("Timed out waiting to release pooled client");
}
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new IllegalStateException(exception);
}
return null;
}));
Assert.assertTrue(borrowed.await(2, TimeUnit.SECONDS));
Future<?> waiter = executor.submit(() -> manager.withClient(client -> null));
try {
waiter.get();
Assert.fail("Pool exhaustion must fail after the configured wait");
} catch (ExecutionException expected) {
Assert.assertNotNull(expected.getCause());
}
release.countDown();
holder.get();
} finally {
release.countDown();
executor.shutdownNow();
}
}
private static Document document(String id, String content, float first, float second) {
Document document = Document.of(content);
document.setId(id);
document.setVector(new float[] { first, second });
return document;
}
}

View File

@@ -5,12 +5,15 @@ import com.easyagents.core.message.SystemMessage;
import com.easyagents.core.model.chat.BaseChatModel;
import com.easyagents.core.model.chat.ChatModel;
import com.easyagents.core.model.chat.StreamResponseListener;
import com.easyagents.core.model.exception.ModelException;
import com.easyagents.core.model.client.StreamContext;
import com.easyagents.core.model.chat.response.AiMessageResponse;
import com.easyagents.core.prompt.SimplePrompt;
import com.easyagents.core.util.ImageUtil;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.WorkflowErrorReason;
import com.easyagents.flow.core.chain.WorkflowExecutionException;
import com.easyagents.flow.core.chain.event.ChainStatusChangeEvent;
import com.easyagents.flow.core.chain.event.LlmStreamEvent;
import com.easyagents.flow.core.chain.listener.ChainEventListener;
@@ -23,14 +26,25 @@ import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.IdentityHashMap;
import java.util.Set;
import java.net.SocketException;
import java.net.SocketTimeoutException;
import java.net.UnknownHostException;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
/**
* 基于 Easy-Agents 聊天模型实现工作流 LLM 调用。
*/
public class EasyAgentsLlm implements Llm {
private static final Pattern MODEL_NOT_FOUND_MESSAGE = Pattern.compile(
"^(?:the\\s+)?model(?:\\s+([`'\"])[^\\r\\n]+\\1)?\\s+(?:not found|does not exist)[.!]?$",
Pattern.CASE_INSENSITIVE);
private ChatModel chatModel;
private ImageInputResolver imageInputResolver;
@@ -163,7 +177,7 @@ public class EasyAgentsLlm implements Llm {
if (message == null || StringUtil.noText(message.getFullContent())) {
failure.compareAndSet(
null,
new IllegalStateException(
new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID,
"EasyAgentsLlm can not get aiMessage!"));
} else {
result.set(message.getFullContent());
@@ -191,6 +205,8 @@ public class EasyAgentsLlm implements Llm {
}
}, chatOptions);
awaitCompletion(completion, streamContext);
} catch (RuntimeException exception) {
throw modelFailure(exception);
} finally {
chain.getEventManager().removeEventListener(
ChainStatusChangeEvent.class, cancellationListener);
@@ -198,14 +214,66 @@ public class EasyAgentsLlm implements Llm {
Throwable throwable = failure.get();
if (throwable != null) {
throw new RuntimeException("EasyAgentsLlm stream failed", throwable);
throw modelFailure(throwable);
}
if (StringUtil.noText(result.get())) {
throw new RuntimeException("EasyAgentsLlm can not get response!");
throw new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID, "EasyAgentsLlm can not get response!");
}
return result.get();
}
static WorkflowExecutionException modelFailure(Throwable error) {
Set<Throwable> seen = Collections.newSetFromMap(new IdentityHashMap<>());
WorkflowErrorReason reason = WorkflowErrorReason.NODE_EXECUTION_FAILED;
for (Throwable cause = error; cause != null && seen.add(cause); cause = cause.getCause()) {
if (cause instanceof WorkflowExecutionException known) {
return known;
}
if (cause instanceof ModelException model && model.getStatusCode() != null) {
int status = model.getStatusCode();
if (status == 429) reason = WorkflowErrorReason.MODEL_RATE_LIMITED;
else if (status == 401 || status == 403) reason = WorkflowErrorReason.MODEL_AUTH_FAILED;
else if (status == 408 || status == 504) reason = WorkflowErrorReason.MODEL_TIMEOUT;
else {
reason = reasonFromModelCode(model.getErrorCode());
if (reason == WorkflowErrorReason.NODE_EXECUTION_FAILED) reason = reasonFromModelCode(model.getErrorType());
if (reason == WorkflowErrorReason.NODE_EXECUTION_FAILED) {
if ((status == 200 || status == 400 || status == 404) && isModelNotFound(model)) {
reason = WorkflowErrorReason.MODEL_NOT_FOUND;
} else if (status >= 500 && status <= 599) {
reason = WorkflowErrorReason.MODEL_UNAVAILABLE;
}
}
}
} else if (cause instanceof SocketTimeoutException || cause instanceof TimeoutException) {
reason = WorkflowErrorReason.MODEL_TIMEOUT;
} else if (cause instanceof SocketException || cause instanceof UnknownHostException) {
reason = WorkflowErrorReason.MODEL_UNAVAILABLE;
}
if (reason != WorkflowErrorReason.NODE_EXECUTION_FAILED) break;
}
return new WorkflowExecutionException(reason, "EasyAgentsLlm stream failed", error);
}
private static WorkflowErrorReason reasonFromModelCode(String code) {
if (code == null) return WorkflowErrorReason.NODE_EXECUTION_FAILED;
return switch (code.toLowerCase(java.util.Locale.ROOT)) {
case "model_not_found" -> WorkflowErrorReason.MODEL_NOT_FOUND;
case "rate_limit_exceeded", "rate_limit_error" -> WorkflowErrorReason.MODEL_RATE_LIMITED;
case "invalid_api_key", "authentication_error", "permission_denied", "permission_error" -> WorkflowErrorReason.MODEL_AUTH_FAILED;
case "service_unavailable", "overloaded_error" -> WorkflowErrorReason.MODEL_UNAVAILABLE;
case "request_timeout", "timeout" -> WorkflowErrorReason.MODEL_TIMEOUT;
default -> WorkflowErrorReason.NODE_EXECUTION_FAILED;
};
}
private static boolean isModelNotFound(ModelException error) {
if ("model_not_found".equalsIgnoreCase(error.getErrorCode())) return true;
String message = error.getErrorMessage();
// 只识别模型服务结构化错误中的明确语义,普通路由 404 不等于模型不存在。
return message != null && message.length() <= 512 && MODEL_NOT_FOUND_MESSAGE.matcher(message.trim()).matches();
}
/**
* 构建模型提示词,并解析图片输入。
*
@@ -292,7 +360,8 @@ public class EasyAgentsLlm implements Llm {
}
String resolvedImage = resolveImage(input);
if (StringUtil.noText(resolvedImage)) {
throw new IllegalArgumentException("Resolved image input must not be blank");
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
"Resolved image input must not be blank");
}
resolvedImages.add(resolvedImage);
}
@@ -315,7 +384,7 @@ public class EasyAgentsLlm implements Llm {
if (imageInput instanceof File file) {
return ImageUtil.imageFileToDataUri(file);
}
throw new IllegalArgumentException(
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
"Unsupported image input type: " + imageInput.getClass().getName());
}
@@ -325,7 +394,8 @@ public class EasyAgentsLlm implements Llm {
private void assertImageSupported() {
if (chatModel instanceof BaseChatModel<?> baseChatModel
&& Boolean.FALSE.equals(baseChatModel.getConfig().getSupportImage())) {
throw new IllegalArgumentException("当前模型不支持图片输入,请选择支持视觉能力的模型");
throw new WorkflowExecutionException(WorkflowErrorReason.INPUT_INVALID,
"当前模型不支持图片输入,请选择支持视觉能力的模型");
}
}
}

View File

@@ -12,6 +12,8 @@ import com.easyagents.core.prompt.Prompt;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.EventManager;
import com.easyagents.flow.core.chain.WorkflowErrorReason;
import com.easyagents.flow.core.chain.WorkflowExecutionException;
import com.easyagents.flow.core.chain.event.LlmStreamEvent;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.node.LlmNode;
@@ -101,14 +103,41 @@ public class EasyAgentsLlmTest {
try {
llm.chat(messageInfo, new Llm.ChatOptions(), null, null);
Assert.fail("expected IllegalArgumentException");
} catch (IllegalArgumentException exception) {
Assert.fail("expected WorkflowExecutionException");
} catch (WorkflowExecutionException exception) {
Assert.assertEquals(WorkflowErrorReason.INPUT_INVALID, exception.getReason());
Assert.assertEquals(
"当前模型不支持图片输入,请选择支持视觉能力的模型",
exception.getMessage());
}
}
@Test
public void shouldClassifyExplicitImageInputValidation() {
EasyAgentsLlm llm = new EasyAgentsLlm();
Llm.MessageInfo message = new Llm.MessageInfo();
message.setImageInputs(List.of(123));
WorkflowExecutionException invalidType = Assert.assertThrows(WorkflowExecutionException.class,
() -> llm.resolveImages(message));
Assert.assertEquals(WorkflowErrorReason.INPUT_INVALID, invalidType.getReason());
llm.setImageInputResolver(input -> " ");
WorkflowExecutionException blank = Assert.assertThrows(WorkflowExecutionException.class,
() -> llm.resolveImages(message));
Assert.assertEquals(WorkflowErrorReason.INPUT_INVALID, blank.getReason());
}
@Test
public void shouldNotReclassifyUnknownImageResolverFailures() {
EasyAgentsLlm llm = new EasyAgentsLlm();
IllegalStateException original = new IllegalStateException("synthetic resolver failure");
llm.setImageInputResolver(input -> { throw original; });
Llm.MessageInfo message = new Llm.MessageInfo();
message.setImageInputs(List.of("image"));
Assert.assertSame(original, Assert.assertThrows(IllegalStateException.class,
() -> llm.resolveImages(message)));
}
/**
* 验证重复执行同一 LLM 节点时,每次调用拥有独立流标识且增量不会被覆盖。
*/

View File

@@ -0,0 +1,45 @@
package com.easyagents.flow.support.provider;
import com.easyagents.core.model.exception.ModelException;
import com.easyagents.flow.core.chain.WorkflowErrorReason;
import com.easyagents.flow.core.chain.WorkflowExecutionException;
import org.junit.Assert;
import org.junit.Test;
import java.net.ConnectException;
import java.net.SocketTimeoutException;
import java.net.SocketException;
public class WorkflowModelFailureTest {
@Test
public void shouldClassifyTypedModelFailuresWithoutParsingMessages() {
assertReason(WorkflowErrorReason.MODEL_RATE_LIMITED, new ModelException(429, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_AUTH_FAILED, new ModelException(401, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_AUTH_FAILED, new ModelException(403, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_UNAVAILABLE, new ModelException(503, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_TIMEOUT, new ModelException(408, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_TIMEOUT, new ModelException(504, "opaque", null));
assertReason(WorkflowErrorReason.MODEL_TIMEOUT, new RuntimeException(new SocketTimeoutException()));
assertReason(WorkflowErrorReason.MODEL_UNAVAILABLE, new ConnectException());
assertReason(WorkflowErrorReason.MODEL_UNAVAILABLE, new SocketException("Connection reset"));
assertReason(WorkflowErrorReason.NODE_EXECUTION_FAILED, new RuntimeException("429 timeout 鉴权失败"));
assertReason(WorkflowErrorReason.NODE_OUTPUT_INVALID, new WorkflowExecutionException(WorkflowErrorReason.NODE_OUTPUT_INVALID, "empty"));
}
@Test
public void shouldRecognizeMissingModelWithoutMisclassifyingRoute404() {
assertReason(WorkflowErrorReason.MODEL_NOT_FOUND,
new ModelException(404, "raw response", null, "404", null, "Model not found"));
assertReason(WorkflowErrorReason.MODEL_NOT_FOUND,
new ModelException(404, "raw response", null, null, null, "The model `test` does not exist."));
assertReason(WorkflowErrorReason.MODEL_NOT_FOUND,
new ModelException(400, "raw response", null, "model_not_found", null, "opaque"));
assertReason(WorkflowErrorReason.NODE_EXECUTION_FAILED,
new ModelException(404, "raw response", null, "404", null, "Not Found"));
assertReason(WorkflowErrorReason.NODE_EXECUTION_FAILED,
new ModelException(404, "raw response", null, "404", null, "Model endpoint not found"));
assertReason(WorkflowErrorReason.MODEL_AUTH_FAILED,
new ModelException(403, "raw response", null, "model_not_found", null, "Model not found"));
}
private void assertReason(WorkflowErrorReason reason, Throwable error) {
Assert.assertEquals(reason, EasyAgentsLlm.modelFailure(error).getReason());
}
}

View File

@@ -0,0 +1,93 @@
package com.easyagents.flow.support.provider;
import com.easyagents.core.model.chat.ChatConfig;
import com.easyagents.core.model.chat.OpenAICompatibleChatModel;
import com.easyagents.flow.core.chain.*;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import com.easyagents.flow.core.llm.LlmManager;
import com.easyagents.flow.core.llm.Llm;
import com.easyagents.flow.core.llm.LlmProvider;
import com.easyagents.flow.core.node.LlmNode;
import com.sun.net.httpserver.HttpServer;
import org.junit.Assert;
import org.junit.Test;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.util.UUID;
/** 真实 HTTP/SSE 到 LlmNode 的分类回归,不直接注入原因码。 */
public class WorkflowModelHttpFailureTest {
@Test
public void shouldClassifyHttpErrorsAtModelBoundary() throws Exception {
assertFailure(429, "{}", false, WorkflowErrorReason.MODEL_RATE_LIMITED);
assertFailure(401, "{}", false, WorkflowErrorReason.MODEL_AUTH_FAILED);
assertFailure(503, "{}", false, WorkflowErrorReason.MODEL_UNAVAILABLE);
assertFailure(408, "{}", false, WorkflowErrorReason.MODEL_TIMEOUT);
assertFailure(504, "{}", false, WorkflowErrorReason.MODEL_TIMEOUT);
assertFailure(404, "{\"error\":{\"message\":\"Model not found\",\"code\":404,\"type\":\"NotFound\"}}",
false, WorkflowErrorReason.MODEL_NOT_FOUND);
assertFailure(404, "<html>Not Found</html>", false, WorkflowErrorReason.NODE_EXECUTION_FAILED);
}
@Test
public void streamErrorAfterPartialOutputMustFailInsteadOfSucceeding() throws Exception {
for (String[] error : new String[][]{
{"rate_limit_exceeded", "MODEL_RATE_LIMITED"},
{"authentication_error", "MODEL_AUTH_FAILED"},
{"overloaded_error", "MODEL_UNAVAILABLE"},
{"request_timeout", "MODEL_TIMEOUT"},
{"model_not_found", "MODEL_NOT_FOUND"}}) {
String body = "data: {\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n"
+ "data: {\"error\":{\"code\":\"upstream_error\",\"type\":\"" + error[0] + "\",\"message\":\"private upstream body\"}}\n\n"
+ "data: [DONE]\n\n";
assertFailure(200, body, false, WorkflowErrorReason.valueOf(error[1]));
}
}
@Test
public void blankOrInvalidJsonOutputMustHaveOutputReason() throws Exception {
assertFailure(200, "data: [DONE]\n\n", false, WorkflowErrorReason.NODE_OUTPUT_INVALID);
for (String output : new String[]{"not-json", "```json\\n\\n```", "null"}) {
assertFailure(200, "data: {\"choices\":[{\"delta\":{\"content\":\"" + output
+ "\"}}]}\n\ndata: [DONE]\n\n", true, WorkflowErrorReason.NODE_OUTPUT_INVALID);
}
}
private void assertFailure(int status, String body, boolean json, WorkflowErrorReason reason) throws Exception {
HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
server.createContext("/chat", exchange -> {
exchange.getRequestBody().readAllBytes();
byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
exchange.getResponseHeaders().set("Content-Type", status == 200 ? "text/event-stream" : "application/json");
exchange.sendResponseHeaders(status, bytes.length);
try (var output = exchange.getResponseBody()) { output.write(bytes); }
});
server.start();
String id = UUID.randomUUID().toString();
ChatConfig config = new ChatConfig();
config.setEndpoint("http://127.0.0.1:" + server.getAddress().getPort());
config.setRequestPath("/chat"); config.setModel("missing-test"); config.setApiKey("synthetic");
config.setLogEnabled(false); config.setObservabilityEnabled(false); config.setRetryEnabled(false);
EasyAgentsLlm llm = new EasyAgentsLlm();
llm.setChatModel(new OpenAICompatibleChatModel<>(config));
LlmProvider provider = modelId -> id.equals(modelId) ? llm : null;
LlmManager.getInstance().registerProvider(provider);
try {
Chain chain = new Chain(new ChainDefinition(), id);
chain.setEventManager(new EventManager());
chain.setChainStateRepository(new InMemoryChainStateRepository());
chain.setNodeStateRepository(new InMemoryNodeStateRepository());
chain.initializeState();
LlmNode node = new LlmNode(); node.setId("llm"); node.setLlmId(id); node.setUserPrompt("test");
node.setChatOptions(new Llm.ChatOptions());
node.setOutType(json ? "json" : "text");
WorkflowExecutionException failure = Assert.assertThrows(WorkflowExecutionException.class, () -> node.execute(chain));
Assert.assertEquals(reason, failure.getReason());
} finally {
LlmManager.getInstance().removeProvider(provider);
server.stop(0);
}
}
}

View File

@@ -54,6 +54,7 @@
<calcite.version>1.42.0</calcite.version>
<h2.version>2.3.232</h2.version>
<quartz.version>2.5.2</quartz.version>
<milvus.version>2.3.11</milvus.version>
</properties>
@@ -152,6 +153,12 @@
<version>${quartz.version}</version>
</dependency>
<dependency>
<groupId>io.milvus</groupId>
<artifactId>milvus-sdk-java</artifactId>
<version>${milvus.version}</version>
</dependency>
<!--easy-agents dependency management-->
<dependency>
<groupId>com.easyagents</groupId>