Compare commits

12 Commits

Author SHA1 Message Date
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
c1fe64cefa perf: 优化多源联邦分片调度
- 增加 Engine 级有界并行调度与资源收口

- 补齐调度饱和、失败快速传播及六源基准测试
2026-09-01 17:02:08 +08:00
1b36067e6c fix: 修复 Lucene 特殊字符查询失败
- 将用户关键词按普通文本转义后再解析

- 避免空查询和解析失败触发二次空指针

- 补充特殊字符与空查询回归测试
2026-09-01 15:11:55 +08:00
876517f821 feat: M28 支持工作流多入边汇聚模式 2026-08-31 15:54:44 +08:00
2d50f7de15 fix: 完善分布式调度恢复与批量触发
- 持久化 Quartz refire 状态并收口运行时启动关闭顺序

- 增加批量 Trigger 获取配置、校验与回归测试
2026-08-31 14:56:41 +08:00
93296eb810 feat: 增强 Agentic RAG 主动检索引导
- 统一组合用户、知识库与异步工具系统提示词

- 补充知识库调用策略与运行时回归测试
2026-08-29 17:01:35 +08:00
47a2706f11 feat: 支持一库一工具 Agentic RAG
- 将知识库声明与 Retriever 绑定为独立 Registration 并注册到 AgentScope Toolkit

- 统一检索分数、最终文档事件与长答案引用语义

- 补齐历史会话、工具类加载器和知识库调用状态回归测试
2026-08-29 15:41:52 +08:00
e146653de7 fix: 保留模型异常后的会话上下文
- 保存模型失败与无正文取消路径中的 AgentScope 会话

- 补充失败和推理中断场景的上下文恢复测试
2026-08-26 22:55:20 +08:00
81 changed files with 8238 additions and 858 deletions

View File

@@ -1,14 +1,17 @@
package com.easyagents.agent.runtime; package com.easyagents.agent.runtime;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetriever; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import com.easyagents.agent.runtime.media.AgentMediaResolver; import com.easyagents.agent.runtime.media.AgentMediaResolver;
import com.easyagents.agent.runtime.memory.AgentMemorySnapshot;
import com.easyagents.agent.runtime.persistence.conversation.AgentConversationRecorder; import com.easyagents.agent.runtime.persistence.conversation.AgentConversationRecorder;
import com.easyagents.agent.runtime.persistence.conversation.noop.NoopAgentConversationRecorder; import com.easyagents.agent.runtime.persistence.conversation.noop.NoopAgentConversationRecorder;
import com.easyagents.agent.runtime.persistence.session.AgentSessionStore; import com.easyagents.agent.runtime.persistence.session.AgentSessionStore;
import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore; import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore;
import com.easyagents.agent.runtime.tool.AgentToolInvoker; import com.easyagents.agent.runtime.tool.AgentToolInvoker;
import java.util.ArrayList;
import java.util.LinkedHashMap; import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map; import java.util.Map;
/** /**
@@ -44,7 +47,12 @@ public class AgentInitRequest {
/** /**
* 知识库集合实现AgentKnowledgeRetriever接口以进行知识检索动作。 * 知识库集合实现AgentKnowledgeRetriever接口以进行知识检索动作。
*/ */
private Map<String, AgentKnowledgeRetriever> knowledgeRetrievers = new LinkedHashMap<>(); private List<AgentKnowledgeRegistration> knowledgeRegistrations = new ArrayList<>();
/**
* 首次构建 Agent 时装载的对话历史快照。
*/
private AgentMemorySnapshot memorySnapshot = new AgentMemorySnapshot();
/** /**
* 对话事件记录器,用于记录运行时事件流。 * 对话事件记录器,用于记录运行时事件流。
@@ -156,17 +164,37 @@ public class AgentInitRequest {
* *
* @return 知识库检索器 * @return 知识库检索器
*/ */
public Map<String, AgentKnowledgeRetriever> getKnowledgeRetrievers() { public List<AgentKnowledgeRegistration> getKnowledgeRegistrations() {
return knowledgeRetrievers; return knowledgeRegistrations;
} }
/** /**
* 设置知识库检索器。 * 设置知识库检索器。
* *
* @param knowledgeRetrievers 知识库检索器 * @param knowledgeRegistrations 知识库运行时绑定
*/ */
public void setKnowledgeRetrievers(Map<String, AgentKnowledgeRetriever> knowledgeRetrievers) { public void setKnowledgeRegistrations(List<AgentKnowledgeRegistration> knowledgeRegistrations) {
this.knowledgeRetrievers = knowledgeRetrievers == null ? new LinkedHashMap<>() : knowledgeRetrievers; this.knowledgeRegistrations = knowledgeRegistrations == null
? new ArrayList<>()
: new ArrayList<>(knowledgeRegistrations);
}
/**
* 获取首次构建 Agent 时的对话历史快照。
*
* @return 对话历史快照
*/
public AgentMemorySnapshot getMemorySnapshot() {
return memorySnapshot;
}
/**
* 设置首次构建 Agent 时的对话历史快照。
*
* @param memorySnapshot 对话历史快照
*/
public void setMemorySnapshot(AgentMemorySnapshot memorySnapshot) {
this.memorySnapshot = memorySnapshot == null ? new AgentMemorySnapshot() : memorySnapshot;
} }
/** /**

View File

@@ -1,6 +1,6 @@
package com.easyagents.agent.runtime; package com.easyagents.agent.runtime;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetriever; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import com.easyagents.agent.runtime.memory.AgentMemorySnapshot; import com.easyagents.agent.runtime.memory.AgentMemorySnapshot;
import com.easyagents.agent.runtime.message.AgentMessage; import com.easyagents.agent.runtime.message.AgentMessage;
import com.easyagents.agent.runtime.persistence.conversation.AgentConversationRecorder; import com.easyagents.agent.runtime.persistence.conversation.AgentConversationRecorder;
@@ -9,7 +9,9 @@ import com.easyagents.agent.runtime.persistence.session.AgentSessionStore;
import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore; import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore;
import com.easyagents.agent.runtime.tool.AgentToolInvoker; import com.easyagents.agent.runtime.tool.AgentToolInvoker;
import java.util.ArrayList;
import java.util.LinkedHashMap; import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map; import java.util.Map;
/** /**
@@ -60,7 +62,7 @@ public class AgentRuntimeExecutionContext {
/** /**
* 按知识库ID索引的检索器。 * 按知识库ID索引的检索器。
*/ */
private Map<String, AgentKnowledgeRetriever> knowledgeRetrievers = new LinkedHashMap<>(); private List<AgentKnowledgeRegistration> knowledgeRegistrations = new ArrayList<>();
/** /**
* 会话状态存储。 * 会话状态存储。
@@ -231,17 +233,19 @@ public class AgentRuntimeExecutionContext {
* *
* @return 知识库检索器 * @return 知识库检索器
*/ */
public Map<String, AgentKnowledgeRetriever> getKnowledgeRetrievers() { public List<AgentKnowledgeRegistration> getKnowledgeRegistrations() {
return knowledgeRetrievers; return knowledgeRegistrations;
} }
/** /**
* 设置知识库检索器。 * 设置知识库检索器。
* *
* @param knowledgeRetrievers 知识库检索器 * @param knowledgeRegistrations 知识库运行时绑定
*/ */
public void setKnowledgeRetrievers(Map<String, AgentKnowledgeRetriever> knowledgeRetrievers) { public void setKnowledgeRegistrations(List<AgentKnowledgeRegistration> knowledgeRegistrations) {
this.knowledgeRetrievers = knowledgeRetrievers == null ? new LinkedHashMap<>() : knowledgeRetrievers; this.knowledgeRegistrations = knowledgeRegistrations == null
? new ArrayList<>()
: new ArrayList<>(knowledgeRegistrations);
} }
/** /**

View File

@@ -2,291 +2,361 @@ package com.easyagents.agent.runtime.agentscope;
import com.easyagents.agent.runtime.AgentRuntimeException; import com.easyagents.agent.runtime.AgentRuntimeException;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext; import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.event.*; import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.knowledge.*; import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import io.agentscope.core.message.TextBlock; import com.easyagents.agent.runtime.event.AgentRuntimeTurnContextHolder;
import io.agentscope.core.rag.Knowledge; import com.easyagents.agent.runtime.hitl.AgentToolApprovalCoordinator;
import io.agentscope.core.rag.model.Document; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeDocument;
import io.agentscope.core.rag.model.DocumentMetadata; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import io.agentscope.core.rag.model.RetrieveConfig; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalRequest;
import reactor.core.publisher.Mono; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalResult;
import reactor.core.publisher.Sinks; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeToolNames;
import com.easyagents.agent.runtime.tool.AgentToolCategory;
import com.easyagents.agent.runtime.tool.AgentToolContext;
import com.easyagents.agent.runtime.tool.AgentToolResult;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.tool.Toolkit;
import java.util.*; import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
/** /**
* 将运行时知识库检索器适配为一个聚合 AgentScope Knowledge * 将中立知识库绑定适配为 AgentScope 一库一工具
*/ */
public class AgentScopeKnowledgeAdapter { public class AgentScopeKnowledgeAdapter {
/** /**
* 创建聚合 Knowledge * 根据 Agent 知识库声明创建模型可见工具定义
* *
* @param request 运行请求 * @param context 运行时上下文
* @return 聚合 Knowledge未配置知识库时返回 null * @return 知识库工具定义
* @throws AgentRuntimeException 声明、运行名或 Retriever 绑定不合法时抛出
*/ */
public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request) { public List<AgentToolSpec> createToolSpecs(AgentRuntimeExecutionContext context) {
return createAggregateKnowledge(request, (Sinks.Many<AgentRuntimeEvent>) null); if (context == null || context.getAgentDefinition() == null) {
throw new AgentRuntimeException("Agent runtime context and definition are required for knowledge tools.");
}
List<AgentKnowledgeSpec> knowledgeSpecs = context.getAgentDefinition().getKnowledgeSpecs();
List<AgentKnowledgeRegistration> registrations = context.getKnowledgeRegistrations();
if (knowledgeSpecs == null || knowledgeSpecs.isEmpty()) {
if (registrations != null && !registrations.isEmpty()) {
throw new AgentRuntimeException("Knowledge registrations require matching knowledge specs.");
}
return List.of();
}
Map<String, AgentKnowledgeRegistration> registrationIndex = registrationIndex(registrations);
if (registrationIndex.size() != knowledgeSpecs.size()) {
throw new AgentRuntimeException("Knowledge specs and registrations must match one-to-one.");
}
List<AgentToolSpec> toolSpecs = new ArrayList<>(knowledgeSpecs.size());
Set<String> toolNames = new LinkedHashSet<>();
for (AgentKnowledgeSpec knowledgeSpec : knowledgeSpecs) {
validateKnowledgeSpec(knowledgeSpec);
AgentKnowledgeRegistration registration = registrationIndex.get(knowledgeSpec.getKnowledgeId());
if (registration == null) {
throw new AgentRuntimeException(
"Knowledge retriever is required: " + knowledgeSpec.getKnowledgeId());
}
String toolName = AgentKnowledgeToolNames.build(knowledgeSpec.getRuntimeName());
if (!toolNames.add(toolName)) {
throw new AgentRuntimeException("Duplicate knowledge tool name: " + toolName);
}
toolSpecs.add(toolSpec(knowledgeSpec, toolName));
}
return toolSpecs;
} }
/** /**
* 创建带事件 sink 的聚合 Knowledge * 将知识库工具注册到现有 AgentScope Toolkit
* *
* @param request 运行请求 * @param context 运行时上下文
* @param eventSink 事件 sink * @param toolSpecs 知识库工具定义
* @return 聚合 Knowledge未配置知识库时返回 null * @param toolkit AgentScope Toolkit
*/ * @param toolAdapter 中立工具适配器
public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request, Sinks.Many<AgentRuntimeEvent> eventSink) { * @param approvalCoordinator 工具审批协调器
return createAggregateKnowledge(request, fixedHolder(request, eventSink));
}
/**
* 创建可读取当前运行轮次事件出口的聚合 Knowledge。
*
* @param request 运行时级上下文
* @param turnContextHolder 当前运行轮次上下文持有器 * @param turnContextHolder 当前运行轮次上下文持有器
* @return 聚合 Knowledge未配置知识库时返回 null
*/ */
public Knowledge createAggregateKnowledge(AgentRuntimeExecutionContext request, public void registerTools(AgentRuntimeExecutionContext context,
List<AgentToolSpec> toolSpecs,
Toolkit toolkit,
AgentScopeToolAdapter toolAdapter,
AgentToolApprovalCoordinator approvalCoordinator,
AgentRuntimeTurnContextHolder turnContextHolder) { AgentRuntimeTurnContextHolder turnContextHolder) {
if (request.getAgentDefinition().getKnowledgeSpecs().isEmpty()) { Objects.requireNonNull(toolkit, "toolkit");
return null; Objects.requireNonNull(toolAdapter, "toolAdapter");
if (toolSpecs == null || toolSpecs.isEmpty()) {
return;
} }
return new AggregateKnowledge(request, turnContextHolder); Map<String, AgentKnowledgeRegistration> registrationIndex =
registrationIndex(context.getKnowledgeRegistrations());
for (AgentToolSpec toolSpec : toolSpecs) {
String knowledgeId = stringValue(toolSpec.getMetadata().get("knowledgeId"));
AgentKnowledgeRegistration registration = registrationIndex.get(knowledgeId);
if (registration == null) {
throw new AgentRuntimeException("Knowledge registration is required: " + knowledgeId);
}
toolkit.registerAgentTool(toolAdapter.adapt(
toolSpec,
(input, toolContext) -> retrieve(registration, input, toolContext),
context,
approvalCoordinator,
turnContextHolder,
null,
null,
false,
false,
false));
} }
private AgentRuntimeTurnContextHolder fixedHolder(AgentRuntimeExecutionContext request,
Sinks.Many<AgentRuntimeEvent> eventSink) {
AgentRuntimeTurnContextHolder holder = new AgentRuntimeTurnContextHolder();
AgentRuntimeEventBridge bridge = new AgentRuntimeEventBridge(request, holder);
holder.set(new AgentRuntimeTurnContext(null, eventSink, bridge));
return holder;
} }
/** /**
* 将运行时文档转换为 AgentScope 文档 * 按知识库 ID 建立运行时绑定索引并拒绝重复绑定
* *
* @param documents 运行时文档 * @param registrations 知识库运行时绑定
* @return AgentScope 文档 * @return 以知识库 ID 为键的绑定索引
*/ */
public List<Document> toDocuments(List<AgentKnowledgeDocument> documents) { private Map<String, AgentKnowledgeRegistration> registrationIndex(
List<Document> converted = new ArrayList<>(); List<AgentKnowledgeRegistration> registrations) {
Map<String, AgentKnowledgeRegistration> index = new LinkedHashMap<>();
if (registrations == null) {
return index;
}
for (AgentKnowledgeRegistration registration : registrations) {
if (registration == null || registration.getKnowledgeSpec() == null) {
throw new AgentRuntimeException("Knowledge registration and spec are required.");
}
String knowledgeId = registration.getKnowledgeSpec().getKnowledgeId();
if (knowledgeId == null || knowledgeId.isBlank()) {
throw new AgentRuntimeException("Knowledge id is required.");
}
if (index.putIfAbsent(knowledgeId, registration) != null) {
throw new AgentRuntimeException("Duplicate knowledge registration: " + knowledgeId);
}
}
return index;
}
/**
* 校验知识库声明的标识和英文运行名。
*
* @param knowledgeSpec 知识库声明
* @throws AgentRuntimeException 声明不合法时抛出
*/
private void validateKnowledgeSpec(AgentKnowledgeSpec knowledgeSpec) {
if (knowledgeSpec == null) {
throw new AgentRuntimeException("Knowledge spec is required.");
}
if (knowledgeSpec.getKnowledgeId() == null || knowledgeSpec.getKnowledgeId().isBlank()) {
throw new AgentRuntimeException("Knowledge id is required.");
}
AgentKnowledgeToolNames.build(knowledgeSpec.getRuntimeName());
}
/**
* 创建单个知识库对应的模型工具声明。
*
* @param knowledgeSpec 知识库声明
* @param toolName 已规范化的工具名
* @return 模型可见工具声明
*/
private AgentToolSpec toolSpec(AgentKnowledgeSpec knowledgeSpec, String toolName) {
AgentToolSpec toolSpec = new AgentToolSpec();
toolSpec.setName(toolName);
toolSpec.setDescription(toolDescription(knowledgeSpec));
toolSpec.setCategory(AgentToolCategory.KNOWLEDGE);
toolSpec.setParametersSchema(Map.of(
"type", "object",
"properties", Map.of(
"query", Map.of(
"type", "string",
"description", "A standalone search query containing all context needed to retrieve relevant knowledge."
)
),
"required", List.of("query"),
"additionalProperties", false
));
toolSpec.getMetadata().putAll(knowledgeSpec.getMetadata());
toolSpec.getMetadata().put("knowledgeId", knowledgeSpec.getKnowledgeId());
toolSpec.getMetadata().put("knowledgeName", knowledgeSpec.getName());
toolSpec.getMetadata().put("knowledgeRuntimeName", knowledgeSpec.getRuntimeName());
toolSpec.getMetadata().put("toolDisplayName", knowledgeSpec.getName());
return toolSpec;
}
/**
* 生成包含知识库名称、范围和查询要求的工具描述。
*
* @param knowledgeSpec 知识库声明
* @return 模型可见工具描述
*/
private String toolDescription(AgentKnowledgeSpec knowledgeSpec) {
String knowledgeName = hasText(knowledgeSpec.getName())
? knowledgeSpec.getName().trim()
: knowledgeSpec.getRuntimeName().trim();
StringBuilder description = new StringBuilder()
.append("Search the knowledge base \"")
.append(knowledgeName)
.append("\" for relevant and grounded information.");
if (hasText(knowledgeSpec.getDescription())) {
description.append(" Knowledge scope: ")
.append(knowledgeSpec.getDescription().trim())
.append('.');
}
description.append(" Build a standalone query from the current question and necessary conversation context.");
return description.toString();
}
/**
* 执行单库检索并生成模型证据及旁路检索事件。
*
* @param registration 知识库运行时绑定
* @param input Function Call 输入
* @param toolContext 工具执行上下文
* @return 模型可见工具结果
*/
private AgentToolResult retrieve(AgentKnowledgeRegistration registration,
Map<String, Object> input,
AgentToolContext toolContext) {
AgentKnowledgeSpec knowledgeSpec = registration.getKnowledgeSpec();
String query = input == null ? null : stringValue(input.get("query"));
if (query == null || query.isBlank()) {
throw new AgentRuntimeException("Knowledge tool query is required: "
+ AgentKnowledgeToolNames.build(knowledgeSpec.getRuntimeName()));
}
String normalizedQuery = query.trim();
AgentKnowledgeRetrievalRequest retrievalRequest = new AgentKnowledgeRetrievalRequest();
retrievalRequest.setQuery(normalizedQuery);
retrievalRequest.setLimit(knowledgeSpec.getLimit());
retrievalRequest.setScoreThreshold(knowledgeSpec.getScoreThreshold());
retrievalRequest.setKnowledgeSpec(knowledgeSpec);
retrievalRequest.setRuntimeContext(toolContext.getRuntimeContext());
retrievalRequest.getMetadata().put("requestId", toolContext.getRequestId());
retrievalRequest.getMetadata().put("traceId", toolContext.getTraceId());
retrievalRequest.getMetadata().put("sessionId", toolContext.getSessionId());
retrievalRequest.getMetadata().put("toolCallId", toolContext.getToolCallId());
AgentKnowledgeRetrievalResult retrievalResult = registration.getRetriever().retrieve(retrievalRequest);
if (retrievalResult == null) {
throw new AgentRuntimeException("Knowledge retriever returned null result: "
+ knowledgeSpec.getKnowledgeId());
}
List<AgentKnowledgeDocument> documents = finalDocuments(knowledgeSpec, retrievalResult.getDocuments());
toolContext.emitEvent(retrievalEvent(toolContext, knowledgeSpec, retrievalRequest, documents));
List<Map<String, Object>> summaries = documentSummaries(documents);
AgentToolResult toolResult = AgentToolResult.success(modelContent(knowledgeSpec, normalizedQuery, documents));
toolResult.setDisplayContent(summaries);
toolResult.getMetadata().putAll(retrievalResult.getMetadata());
toolResult.getMetadata().put("knowledgeId", knowledgeSpec.getKnowledgeId());
toolResult.getMetadata().put("knowledgeName", knowledgeSpec.getName());
toolResult.getMetadata().put("query", normalizedQuery);
toolResult.getMetadata().put("documentCount", documents.size());
toolResult.getMetadata().put("documents", summaries);
return toolResult;
}
/**
* 对检索结果执行统一阈值过滤、降序排序和绑定数量截断。
*
* @param knowledgeSpec 知识库声明
* @param documents 原始检索文档
* @return 最终进入模型上下文的文档
*/
private List<AgentKnowledgeDocument> finalDocuments(AgentKnowledgeSpec knowledgeSpec,
List<AgentKnowledgeDocument> documents) {
List<AgentKnowledgeDocument> finalDocuments = new ArrayList<>();
if (documents == null) { if (documents == null) {
return converted; return finalDocuments;
} }
for (AgentKnowledgeDocument document : documents) { for (AgentKnowledgeDocument document : documents) {
converted.add(toDocument(document)); if (document == null || !passesThreshold(document, knowledgeSpec.getScoreThreshold())) {
}
return converted;
}
private Document toDocument(AgentKnowledgeDocument document) {
Map<String, Object> payload = new LinkedHashMap<>();
payload.put("documentId", document.getDocumentId());
payload.put("documentName", document.getDocumentName());
payload.put("chunkId", document.getChunkId());
payload.put("sourceUri", document.getSourceUri());
payload.put("knowledgeMetadata", document.getKnowledgeMetadata());
payload.put("documentMetadata", document.getMetadata());
payload.putAll(document.getMetadata());
DocumentMetadata metadata = DocumentMetadata.builder()
.content(TextBlock.builder().text(safeContent(document)).build())
.docId(safeDocumentId(document))
.chunkId(safeChunkId(document))
.payload(payload)
.build();
Document converted = new Document(metadata);
converted.setScore(document.getScore());
return converted;
}
/**
* 获取 AgentScope 要求的非空文档 ID。
*
* @param document 知识文档
* @return 非空文档 ID
*/
private String safeDocumentId(AgentKnowledgeDocument document) {
if (document.getDocumentId() != null && !document.getDocumentId().isBlank()) {
return document.getDocumentId();
}
if (document.getChunkId() != null && !document.getChunkId().isBlank()) {
return document.getChunkId();
}
return "knowledge-document";
}
/**
* 获取 AgentScope 要求的非空分片 ID。
*
* @param document 知识文档
* @return 非空分片 ID
*/
private String safeChunkId(AgentKnowledgeDocument document) {
if (document.getChunkId() != null && !document.getChunkId().isBlank()) {
return document.getChunkId();
}
if (document.getDocumentId() != null && !document.getDocumentId().isBlank()) {
return document.getDocumentId();
}
return "0";
}
/**
* 获取 AgentScope 要求的非空文档内容。
*
* @param document 知识文档
* @return 文档内容
*/
private String safeContent(AgentKnowledgeDocument document) {
return document.getContent() == null ? "" : document.getContent();
}
/**
* 将检索调用分发到多个知识源的聚合 Knowledge 实现。
*/
private class AggregateKnowledge implements Knowledge {
private final AgentRuntimeExecutionContext request;
private final AgentRuntimeTurnContextHolder turnContextHolder;
private AggregateKnowledge(AgentRuntimeExecutionContext request, AgentRuntimeTurnContextHolder turnContextHolder) {
this.request = request;
this.turnContextHolder = turnContextHolder;
}
/**
* 忽略文档新增,因为知识库索引由 EasyFlow 负责。
*
* @param documents 文档列表
* @return 完成信号
*/
@Override
public Mono<Void> addDocuments(List<Document> documents) {
return Mono.error(new UnsupportedOperationException(
"Easy-Agents agent runtime knowledge does not support addDocuments. Use external knowledge service instead."));
}
/**
* 从已配置的知识源检索文档。
*
* @param query 查询
* @param config 检索配置
* @return 文档列表
*/
@Override
public Mono<List<Document>> retrieve(String query, RetrieveConfig config) {
return Mono.fromCallable(() -> retrieveAll(query, config));
}
/**
* 检索并合并所有已配置的知识源。
*
* @param query 查询
* @param config 检索配置
* @return 合并后的文档
*/
private List<Document> retrieveAll(String query, RetrieveConfig config) {
List<AgentKnowledgeDocument> allDocuments = new ArrayList<>();
int globalLimit = config == null || config.getLimit() <= 0 ? 5 : config.getLimit();
double globalThreshold = config == null ? 0D : config.getScoreThreshold();
for (AgentKnowledgeSpec spec : request.getAgentDefinition().getKnowledgeSpecs()) {
AgentKnowledgeRetriever retriever = request.getKnowledgeRetrievers().get(spec.getKnowledgeId());
if (retriever == null) {
throw new AgentRuntimeException("Knowledge retriever is required: " + spec.getKnowledgeId());
}
AgentKnowledgeRetrievalRequest retrievalRequest = new AgentKnowledgeRetrievalRequest();
retrievalRequest.setQuery(query);
retrievalRequest.setLimit(spec.getLimit());
retrievalRequest.setScoreThreshold(Math.max(spec.getScoreThreshold(), globalThreshold));
retrievalRequest.setKnowledgeSpec(spec);
AgentRuntimeExecutionContext currentRequest = currentRequest();
retrievalRequest.setRuntimeContext(currentRequest.getRuntimeContext());
retrievalRequest.getMetadata().put("traceId", currentRequest.getTraceId());
retrievalRequest.getMetadata().put("sessionId", currentRequest.getSessionId());
AgentKnowledgeRetrievalResult result = retriever.retrieve(retrievalRequest);
if (result == null || result.getDocuments() == null) {
emitKnowledgeRetrievalEvent(query, spec, retrievalRequest, new ArrayList<>());
continue; continue;
} }
emitKnowledgeRetrievalEvent(query, spec, retrievalRequest, result.getDocuments()); preserveKnowledgeMetadata(knowledgeSpec, document);
for (AgentKnowledgeDocument document : result.getDocuments()) { finalDocuments.add(document);
preserveKnowledgeMetadata(spec, document);
allDocuments.add(document);
} }
} finalDocuments.sort(Comparator.comparing(
allDocuments.sort(Comparator.comparing(
AgentKnowledgeDocument::getScore, AgentKnowledgeDocument::getScore,
Comparator.nullsLast(Comparator.reverseOrder()) Comparator.nullsLast(Comparator.reverseOrder())));
)); int limit = Math.max(knowledgeSpec.getLimit(), 1);
if (allDocuments.size() > globalLimit) { if (finalDocuments.size() > limit) {
allDocuments = new ArrayList<>(allDocuments.subList(0, globalLimit)); return new ArrayList<>(finalDocuments.subList(0, limit));
} }
return toDocuments(allDocuments); return finalDocuments;
} }
/** /**
* 在单条文档上保留知识库级元数据 * 判断文档最终分数是否达到绑定阈值
* *
* @param spec 知识库声明 * @param document 检索文档
* @param document 文档 * @param scoreThreshold 分数阈值
* @return 达到阈值时为 true
*/ */
private void preserveKnowledgeMetadata(AgentKnowledgeSpec spec, AgentKnowledgeDocument document) { private boolean passesThreshold(AgentKnowledgeDocument document, double scoreThreshold) {
Map<String, Object> knowledgeMetadata = new LinkedHashMap<>(spec.getMetadata()); if (scoreThreshold <= 0D) {
knowledgeMetadata.put("knowledgeId", spec.getKnowledgeId()); return true;
knowledgeMetadata.put("knowledgeName", spec.getName()); }
knowledgeMetadata.put("retrievalMode", spec.getRetrievalMode().name()); return document.getScore() != null && document.getScore() >= scoreThreshold;
}
/**
* 将知识库归属信息合并到文档元数据中。
*
* @param knowledgeSpec 知识库声明
* @param document 检索文档
*/
private void preserveKnowledgeMetadata(AgentKnowledgeSpec knowledgeSpec,
AgentKnowledgeDocument document) {
Map<String, Object> knowledgeMetadata = new LinkedHashMap<>(knowledgeSpec.getMetadata());
knowledgeMetadata.put("knowledgeId", knowledgeSpec.getKnowledgeId());
knowledgeMetadata.put("knowledgeName", knowledgeSpec.getName());
knowledgeMetadata.put("knowledgeRuntimeName", knowledgeSpec.getRuntimeName());
knowledgeMetadata.putAll(document.getKnowledgeMetadata()); knowledgeMetadata.putAll(document.getKnowledgeMetadata());
document.setKnowledgeMetadata(knowledgeMetadata); document.setKnowledgeMetadata(knowledgeMetadata);
document.getMetadata().putIfAbsent("knowledgeId", spec.getKnowledgeId()); document.getMetadata().putIfAbsent("knowledgeId", knowledgeSpec.getKnowledgeId());
document.getMetadata().putIfAbsent("knowledgeName", spec.getName()); document.getMetadata().putIfAbsent("knowledgeName", knowledgeSpec.getName());
} }
/** /**
* 发射知识库检索旁路事件,供聊天界面展示检索过程 * 创建与模型最终证据一致的知识库检索事件
* *
* <p>知识库检索本身属于 AgentScope RAG 主线路,返回的 Document 会继续进入 * @param toolContext 工具执行上下文
* AgentScope 的上下文注入流程;这里发出的 {@code KNOWLEDGE_RETRIEVAL} * @param knowledgeSpec 知识库声明
* 只是旁路告知调用方,不会回写 memory也不会参与模型消息序列。</p>
*
* @param query 查询
* @param spec 知识库声明
* @param retrievalRequest 检索请求 * @param retrievalRequest 检索请求
* @param documents 检索文档 * @param documents 最终文档
* @return 检索旁路事件
*/ */
private void emitKnowledgeRetrievalEvent(String query, private AgentRuntimeEvent retrievalEvent(AgentToolContext toolContext,
AgentKnowledgeSpec spec, AgentKnowledgeSpec knowledgeSpec,
AgentKnowledgeRetrievalRequest retrievalRequest, AgentKnowledgeRetrievalRequest retrievalRequest,
List<AgentKnowledgeDocument> documents) { List<AgentKnowledgeDocument> documents) {
AgentRuntimeEvent event = currentEventBridge().event(AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL); AgentRuntimeEvent event = AgentRuntimeEvent.of(AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL);
event.getPayload().put("query", query); event.setToolCallId(toolContext.getToolCallId());
event.getPayload().put("knowledgeId", spec.getKnowledgeId()); event.getPayload().put("query", retrievalRequest.getQuery());
event.getPayload().put("knowledgeName", spec.getName()); event.getPayload().put("knowledgeId", knowledgeSpec.getKnowledgeId());
event.getPayload().put("knowledgeType", spec.getMetadata().get("knowledgeType")); event.getPayload().put("knowledgeName", knowledgeSpec.getName());
event.getPayload().put("faqCollection", spec.getMetadata().get("faqCollection")); event.getPayload().put("knowledgeType", knowledgeSpec.getMetadata().get("knowledgeType"));
event.getPayload().put("faqCollection", knowledgeSpec.getMetadata().get("faqCollection"));
event.getPayload().put("limit", retrievalRequest.getLimit()); event.getPayload().put("limit", retrievalRequest.getLimit());
event.getPayload().put("scoreThreshold", retrievalRequest.getScoreThreshold()); event.getPayload().put("scoreThreshold", retrievalRequest.getScoreThreshold());
event.getPayload().put("documentCount", documents == null ? 0 : documents.size()); event.getPayload().put("documentCount", documents.size());
event.getPayload().put("documents", documentSummaries(documents)); event.getPayload().put("documents", documentSummaries(documents));
currentEventBridge().emit(event); event.getMetadata().put("toolName", AgentKnowledgeToolNames.build(knowledgeSpec.getRuntimeName()));
} event.getMetadata().put("knowledgeRuntimeName", knowledgeSpec.getRuntimeName());
return event;
private AgentRuntimeExecutionContext currentRequest() {
return turnContextHolder == null ? request : turnContextHolder.executionContext(request);
}
private AgentRuntimeEventBridge currentEventBridge() {
if (turnContextHolder != null && turnContextHolder.eventBridge().isPresent()) {
return turnContextHolder.eventBridge().get();
}
return new AgentRuntimeEventBridge(request, turnContextHolder);
} }
/** /**
* 构建用于事件展示的命中片段,保留前端引注需要的原始 chunk 内容 * 将最终文档转换为 UI 和完成消息引用可消费的稳定摘要
* *
* @param documents 检索文档 * @param documents 最终文档
* @return 命中片段列表 * @return 文档摘要列表
*/ */
private List<Map<String, Object>> documentSummaries(List<AgentKnowledgeDocument> documents) { private List<Map<String, Object>> documentSummaries(List<AgentKnowledgeDocument> documents) {
List<Map<String, Object>> summaries = new ArrayList<>(); List<Map<String, Object>> summaries = new ArrayList<>();
@@ -306,5 +376,78 @@ public class AgentScopeKnowledgeAdapter {
} }
return summaries; return summaries;
} }
/**
* 格式化模型可见的结构化检索证据。
*
* @param knowledgeSpec 知识库声明
* @param query 实际检索词
* @param documents 最终文档
* @return 模型上下文文本
*/
private String modelContent(AgentKnowledgeSpec knowledgeSpec,
String query,
List<AgentKnowledgeDocument> documents) {
String knowledgeName = hasText(knowledgeSpec.getName())
? knowledgeSpec.getName().trim()
: knowledgeSpec.getRuntimeName().trim();
if (documents.isEmpty()) {
return "No relevant documents were found in knowledge base \""
+ knowledgeName + "\" for query: " + query;
}
StringBuilder content = new StringBuilder()
.append("Retrieved evidence from knowledge base \"")
.append(knowledgeName)
.append("\" for query: ")
.append(query)
.append("\n\n");
for (int index = 0; index < documents.size(); index++) {
AgentKnowledgeDocument document = documents.get(index);
content.append("[Evidence ").append(index + 1).append("]\n");
appendField(content, "Document", document.getDocumentName());
appendField(content, "Document ID", document.getDocumentId());
appendField(content, "Chunk ID", document.getChunkId());
appendField(content, "Source", document.getSourceUri());
if (document.getScore() != null) {
content.append("Score: ").append(document.getScore()).append('\n');
}
content.append("Content:\n")
.append(document.getContent() == null ? "" : document.getContent())
.append("\n\n");
}
return content.toString().stripTrailing();
}
/**
* 追加非空证据字段。
*
* @param content 输出缓冲区
* @param label 字段标签
* @param value 字段值
*/
private void appendField(StringBuilder content, String label, String value) {
if (hasText(value)) {
content.append(label).append(": ").append(value.trim()).append('\n');
}
}
/**
* 判断字符串是否包含非空白文本。
*
* @param value 待判断字符串
* @return 包含文本时为 true
*/
private boolean hasText(String value) {
return value != null && !value.isBlank();
}
/**
* 将可空值转换为字符串。
*
* @param value 原始值
* @return 字符串;原始值为空时返回 null
*/
private String stringValue(Object value) {
return value == null ? null : String.valueOf(value);
} }
} }

View File

@@ -13,7 +13,6 @@ import com.easyagents.agent.runtime.hitl.AgentPendingState;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalCoordinator; import com.easyagents.agent.runtime.hitl.AgentToolApprovalCoordinator;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalResolution; import com.easyagents.agent.runtime.hitl.AgentToolApprovalResolution;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalRejectedException; import com.easyagents.agent.runtime.hitl.AgentToolApprovalRejectedException;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.knowledge.citation.AgentKnowledgeCitationMatcher; import com.easyagents.agent.runtime.knowledge.citation.AgentKnowledgeCitationMatcher;
import com.easyagents.agent.runtime.knowledge.citation.HeuristicKnowledgeCitationMatcher; import com.easyagents.agent.runtime.knowledge.citation.HeuristicKnowledgeCitationMatcher;
import com.easyagents.agent.runtime.message.*; import com.easyagents.agent.runtime.message.*;
@@ -22,6 +21,7 @@ import com.easyagents.agent.runtime.mcp.McpSkillRegistration;
import com.easyagents.agent.runtime.mcp.McpSpecValidator; import com.easyagents.agent.runtime.mcp.McpSpecValidator;
import com.easyagents.agent.runtime.mcp.McpToolkitAdapter; import com.easyagents.agent.runtime.mcp.McpToolkitAdapter;
import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore; import com.easyagents.agent.runtime.persistence.session.noop.NoopAgentSessionStore;
import com.easyagents.agent.runtime.prompt.SystemPromptComposer;
import com.easyagents.agent.runtime.skill.AgentSkillBinding; import com.easyagents.agent.runtime.skill.AgentSkillBinding;
import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext; import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext;
import com.easyagents.agent.runtime.tool.AgentToolInvoker; import com.easyagents.agent.runtime.tool.AgentToolInvoker;
@@ -35,9 +35,6 @@ import io.agentscope.core.memory.Memory;
import io.agentscope.core.memory.autocontext.AutoContextMemory; import io.agentscope.core.memory.autocontext.AutoContextMemory;
import io.agentscope.core.message.*; import io.agentscope.core.message.*;
import io.agentscope.core.model.Model; import io.agentscope.core.model.Model;
import io.agentscope.core.rag.Knowledge;
import io.agentscope.core.rag.RAGMode;
import io.agentscope.core.rag.model.RetrieveConfig;
import io.agentscope.core.session.Session; import io.agentscope.core.session.Session;
import io.agentscope.core.skill.SkillBox; import io.agentscope.core.skill.SkillBox;
import io.agentscope.core.state.SessionKey; import io.agentscope.core.state.SessionKey;
@@ -58,18 +55,6 @@ import java.util.function.Supplier;
*/ */
public class AgentScopeReActRuntime implements AgentRuntime { public class AgentScopeReActRuntime implements AgentRuntime {
private static final String ASYNC_TOOL_SYSTEM_PROMPT = """
Async tool protocol:
- Async tools may expose submit, observe, result, cancel, and list sub-tools. Treat these sub-tools as one user-facing tool.
- Do not ask the user to choose submit, observe, result, cancel, or list. These are internal execution phases.
- For a normal user request to use an async tool, call its submit sub-tool first with the user-provided arguments by default.
- After submit returns task_id, immediately call observe with that task_id to check progress.
- If the task is completed and result is available, use the returned result to answer the user.
- If the task is still running after observation, tell the user that the task is running and keep task_id/next_action for later tool calls.
- Use result, list, or cancel directly only when the user explicitly asks to get a known task result, list tasks, or cancel a task.
""";
private final AgentScopeModelFactory modelFactory; private final AgentScopeModelFactory modelFactory;
private final AgentScopeToolAdapter toolAdapter; private final AgentScopeToolAdapter toolAdapter;
private final AgentScopeKnowledgeAdapter knowledgeAdapter; private final AgentScopeKnowledgeAdapter knowledgeAdapter;
@@ -360,7 +345,7 @@ public class AgentScopeReActRuntime implements AgentRuntime {
context.setRuntimeContext(runtimeContext.getRuntimeContext()); context.setRuntimeContext(runtimeContext.getRuntimeContext());
context.setUserMessage(userMessage); context.setUserMessage(userMessage);
context.setToolInvokers(runtimeContext.getToolInvokers()); context.setToolInvokers(runtimeContext.getToolInvokers());
context.setKnowledgeRetrievers(runtimeContext.getKnowledgeRetrievers()); context.setKnowledgeRegistrations(runtimeContext.getKnowledgeRegistrations());
context.setSessionStore(runtimeContext.getSessionStore()); context.setSessionStore(runtimeContext.getSessionStore());
context.setConversationRecorder(runtimeContext.getConversationRecorder()); context.setConversationRecorder(runtimeContext.getConversationRecorder());
context.setMetadata(runtimeContext.getMetadata()); context.setMetadata(runtimeContext.getMetadata());
@@ -381,7 +366,7 @@ public class AgentScopeReActRuntime implements AgentRuntime {
context.setAgentDefinition(runtimeContext.getAgentDefinition()); context.setAgentDefinition(runtimeContext.getAgentDefinition());
context.setRuntimeContext(runtimeContext.getRuntimeContext()); context.setRuntimeContext(runtimeContext.getRuntimeContext());
context.setToolInvokers(runtimeContext.getToolInvokers()); context.setToolInvokers(runtimeContext.getToolInvokers());
context.setKnowledgeRetrievers(runtimeContext.getKnowledgeRetrievers()); context.setKnowledgeRegistrations(runtimeContext.getKnowledgeRegistrations());
context.setSessionStore(runtimeContext.getSessionStore()); context.setSessionStore(runtimeContext.getSessionStore());
context.setConversationRecorder(runtimeContext.getConversationRecorder()); context.setConversationRecorder(runtimeContext.getConversationRecorder());
Map<String, Object> metadata = new LinkedHashMap<>(runtimeContext.getMetadata()); Map<String, Object> metadata = new LinkedHashMap<>(runtimeContext.getMetadata());
@@ -689,6 +674,7 @@ public class AgentScopeReActRuntime implements AgentRuntime {
saveSession(); saveSession();
return Flux.just(cancelled(context)); return Flux.just(cancelled(context));
} }
saveSession();
return Flux.just(failed(context, error)); return Flux.just(failed(context, error));
} }
@@ -733,7 +719,8 @@ public class AgentScopeReActRuntime implements AgentRuntime {
* 将取消前已输出的助手内容补写入 AgentScope memory 并保存 session。 * 将取消前已输出的助手内容补写入 AgentScope memory 并保存 session。
* *
* <p>AgentScope 的正常完成路径会自行把最终助手消息写入 memory。取消订阅时不会触发 * <p>AgentScope 的正常完成路径会自行把最终助手消息写入 memory。取消订阅时不会触发
* 完成路径,因此这里仅在已有非空助手内容时补写一次,确保下一轮对话能拿到中断前上下文。</p> * 完成路径,因此这里仅在已有非空助手内容时补写一次,并始终保存已经进入 memory 的
* 用户消息,确保下一轮对话能拿到中断前上下文。</p>
* *
* @param finalText 当前已累计的助手文本 * @param finalText 当前已累计的助手文本
* @param finalMessage 当前已捕获的结构化助手消息 * @param finalMessage 当前已捕获的结构化助手消息
@@ -741,10 +728,9 @@ public class AgentScopeReActRuntime implements AgentRuntime {
private void persistPartialAssistantOnCancel(StringBuilder finalText, private void persistPartialAssistantOnCancel(StringBuilder finalText,
AtomicReference<AgentMessage> finalMessage) { AtomicReference<AgentMessage> finalMessage) {
AgentMessage partialMessage = partialAssistantMessage(finalText, finalMessage); AgentMessage partialMessage = partialAssistantMessage(finalText, finalMessage);
if (partialMessage == null) { if (partialMessage != null) {
return;
}
agent.getMemory().addMessage(messageAdapter.toMsg(partialMessage)); agent.getMemory().addMessage(messageAdapter.toMsg(partialMessage));
}
saveSession(); saveSession();
} }
@@ -1109,7 +1095,8 @@ public class AgentScopeReActRuntime implements AgentRuntime {
context.setAgentDefinition(request.getAgentDefinition()); context.setAgentDefinition(request.getAgentDefinition());
context.setRuntimeContext(request.getRuntimeContext()); context.setRuntimeContext(request.getRuntimeContext());
context.setToolInvokers(request.getToolInvokers()); context.setToolInvokers(request.getToolInvokers());
context.setKnowledgeRetrievers(request.getKnowledgeRetrievers()); context.setKnowledgeRegistrations(request.getKnowledgeRegistrations());
context.setMemorySnapshot(request.getMemorySnapshot());
context.setSessionStore(request.getSessionStore()); context.setSessionStore(request.getSessionStore());
context.setConversationRecorder(request.getConversationRecorder()); context.setConversationRecorder(request.getConversationRecorder());
context.setMetadata(request.getMetadata()); context.setMetadata(request.getMetadata());
@@ -1128,9 +1115,9 @@ public class AgentScopeReActRuntime implements AgentRuntime {
Toolkit toolkit = new Toolkit(); Toolkit toolkit = new Toolkit();
AgentScopeToolkitBuildResult toolkitBuildResult = buildToolkit(context, toolkit); AgentScopeToolkitBuildResult toolkitBuildResult = buildToolkit(context, toolkit);
Map<String, List<AgentTool>> skillTools = toolkitBuildResult.skillTools(); Map<String, List<AgentTool>> skillTools = toolkitBuildResult.skillTools();
AgentScopeMemoryBuildResult memoryResult = memoryAdapter.createMemoryResult(null, definition.getMemoryPolicy(), model); AgentScopeMemoryBuildResult memoryResult = memoryAdapter.createMemoryResult(
context.getMemorySnapshot(), definition.getMemoryPolicy(), model);
Memory memory = memoryResult.getMemory(); Memory memory = memoryResult.getMemory();
Knowledge knowledge = knowledgeAdapter.createAggregateKnowledge(context, turnContextHolder);
SkillBox skillBox = skillAdapter.createSkillBox(definition.getSkillBoxSpec(), toolkit, skillTools, SkillBox skillBox = skillAdapter.createSkillBox(definition.getSkillBoxSpec(), toolkit, skillTools,
toolkitBuildResult.skillMcpRegistrations()); toolkitBuildResult.skillMcpRegistrations());
// AutoContextInterceptor 是官方 AutoContextHook 的替代实现。这里仍只注册统一 runtime hook // AutoContextInterceptor 是官方 AutoContextHook 的替代实现。这里仍只注册统一 runtime hook
@@ -1141,7 +1128,10 @@ public class AgentScopeReActRuntime implements AgentRuntime {
interceptors.add(new AutoContextInterceptor(eventBridge, memoryResult.getAutoContextConfig())); interceptors.add(new AutoContextInterceptor(eventBridge, memoryResult.getAutoContextConfig()));
} }
interceptors.add(new MediaReferenceInterceptor(initRequest.getMediaResolver())); interceptors.add(new MediaReferenceInterceptor(initRequest.getMediaResolver()));
List<AgentToolSpec> runtimeToolSpecs = mergeToolSpecs(definition.getToolSpecs(), toolkitBuildResult.mcpToolSpecs(), List<AgentToolSpec> runtimeToolSpecs = mergeToolSpecs(
definition.getToolSpecs(),
toolkitBuildResult.knowledgeToolSpecs(),
toolkitBuildResult.mcpToolSpecs(),
toolkitBuildResult.operateToolSpecs()); toolkitBuildResult.operateToolSpecs());
interceptors.add(new ToolHitlInterceptor(eventBridge, approvalCoordinator, interceptors.add(new ToolHitlInterceptor(eventBridge, approvalCoordinator,
runtimeToolSpecs)); runtimeToolSpecs));
@@ -1156,7 +1146,7 @@ public class AgentScopeReActRuntime implements AgentRuntime {
ReActAgent.Builder builder = ReActAgent.builder() ReActAgent.Builder builder = ReActAgent.builder()
.name(definition.getAgentName()) .name(definition.getAgentName())
.description(definition.getDescription()) .description(definition.getDescription())
.sysPrompt(systemPrompt(definition)) .sysPrompt(SystemPromptComposer.compose(definition))
.model(model) .model(model)
.toolkit(toolkit) .toolkit(toolkit)
.memory(memory) .memory(memory)
@@ -1165,41 +1155,12 @@ public class AgentScopeReActRuntime implements AgentRuntime {
.hook(new AgentScopeRuntimeHook(observationManager)) .hook(new AgentScopeRuntimeHook(observationManager))
.enablePendingToolRecovery(true) .enablePendingToolRecovery(true)
.statePersistence(AgentScopeSessionAdapter.toStatePersistence(definition.getPersistencePolicy())); .statePersistence(AgentScopeSessionAdapter.toStatePersistence(definition.getPersistencePolicy()));
if (knowledge != null) {
builder.knowledge(knowledge)
.ragMode(RAGMode.AGENTIC)
.retrieveConfig(defaultRetrieveConfig(definition));
}
if (skillBox != null) { if (skillBox != null) {
builder.skillBox(skillBox); builder.skillBox(skillBox);
} }
return builder.build(); return builder.build();
} }
private String systemPrompt(AgentDefinition definition) {
String prompt = definition.getSystemPrompt();
if (!hasAsyncTool(definition)) {
return prompt;
}
if (prompt == null || prompt.isBlank()) {
return ASYNC_TOOL_SYSTEM_PROMPT.strip();
}
return prompt.stripTrailing() + ASYNC_TOOL_SYSTEM_PROMPT;
}
private boolean hasAsyncTool(AgentDefinition definition) {
if (definition == null || definition.getToolSpecs() == null) {
return false;
}
for (AgentToolSpec toolSpec : definition.getToolSpecs()) {
// AsyncToolSpecExpander marks all generated sub-tools with this runtime metadata.
if (toolSpec != null && Boolean.TRUE.equals(toolSpec.getMetadata().get("asyncTool"))) {
return true;
}
}
return false;
}
/** /**
* 构建 AgentScope Toolkit并返回按 Skill ID 分组的工具。 * 构建 AgentScope Toolkit并返回按 Skill ID 分组的工具。
* *
@@ -1211,8 +1172,11 @@ public class AgentScopeReActRuntime implements AgentRuntime {
Toolkit toolkit) { Toolkit toolkit) {
Map<String, List<AgentTool>> skillTools = new LinkedHashMap<>(); Map<String, List<AgentTool>> skillTools = new LinkedHashMap<>();
if (!context.getAgentDefinition().getExecutionOptions().isToolCallingEnabled()) { if (!context.getAgentDefinition().getExecutionOptions().isToolCallingEnabled()) {
return new AgentScopeToolkitBuildResult(skillTools, List.of(), List.of(), List.of()); return new AgentScopeToolkitBuildResult(skillTools, List.of(), List.of(), List.of(), List.of());
} }
List<AgentToolSpec> knowledgeToolSpecs = knowledgeAdapter.createToolSpecs(context);
validateRuntimeToolConflicts(context.getAgentDefinition().getToolSpecs(), knowledgeToolSpecs,
List.of(), List.of());
for (AgentToolSpec toolSpec : context.getAgentDefinition().getToolSpecs()) { for (AgentToolSpec toolSpec : context.getAgentDefinition().getToolSpecs()) {
AgentToolInvoker invoker = context.getToolInvokers().get(toolSpec.getName()); AgentToolInvoker invoker = context.getToolInvokers().get(toolSpec.getName());
AgentSkillBinding skillBinding = skillContext.getToolBinding(toolSpec.getName()); AgentSkillBinding skillBinding = skillContext.getToolBinding(toolSpec.getName());
@@ -1224,6 +1188,8 @@ public class AgentScopeReActRuntime implements AgentRuntime {
skillTools.computeIfAbsent(skillBinding.getSkillId(), key -> new ArrayList<>()).add(agentTool); skillTools.computeIfAbsent(skillBinding.getSkillId(), key -> new ArrayList<>()).add(agentTool);
} }
} }
knowledgeAdapter.registerTools(context, knowledgeToolSpecs, toolkit, toolAdapter,
approvalCoordinator, turnContextHolder);
McpRegistration mcpRegistration = mcpToolkitAdapter.register( McpRegistration mcpRegistration = mcpToolkitAdapter.register(
context.getAgentDefinition().getMcpSpecs(), toolkit); context.getAgentDefinition().getMcpSpecs(), toolkit);
mcpClients.addAll(mcpRegistration.getClients()); mcpClients.addAll(mcpRegistration.getClients());
@@ -1231,17 +1197,32 @@ public class AgentScopeReActRuntime implements AgentRuntime {
context.getAgentDefinition().getOperateToolSpecs(), toolkit); context.getAgentDefinition().getOperateToolSpecs(), toolkit);
McpSpecValidator.validateToolConflicts(context.getAgentDefinition().getToolSpecs(), McpSpecValidator.validateToolConflicts(context.getAgentDefinition().getToolSpecs(),
mcpRegistration.getToolSpecs(), context.getAgentDefinition().getOperateToolSpecs()); mcpRegistration.getToolSpecs(), context.getAgentDefinition().getOperateToolSpecs());
return new AgentScopeToolkitBuildResult(skillTools, mcpRegistration.getToolSpecs(), operateToolSpecs, validateRuntimeToolConflicts(context.getAgentDefinition().getToolSpecs(), knowledgeToolSpecs,
mcpRegistration.getSkillRegistrations()); mcpRegistration.getToolSpecs(), operateToolSpecs);
return new AgentScopeToolkitBuildResult(skillTools, knowledgeToolSpecs,
mcpRegistration.getToolSpecs(), operateToolSpecs, mcpRegistration.getSkillRegistrations());
} }
/**
* 合并所有运行时工具声明供统一治理与事件展示使用。
*
* @param toolSpecs 普通工具声明
* @param knowledgeToolSpecs 知识库工具声明
* @param mcpToolSpecs MCP 工具声明
* @param operateToolSpecs 操作工具声明
* @return 保持注册顺序的工具声明列表
*/
private List<AgentToolSpec> mergeToolSpecs(List<AgentToolSpec> toolSpecs, private List<AgentToolSpec> mergeToolSpecs(List<AgentToolSpec> toolSpecs,
List<AgentToolSpec> knowledgeToolSpecs,
List<AgentToolSpec> mcpToolSpecs, List<AgentToolSpec> mcpToolSpecs,
List<AgentToolSpec> operateToolSpecs) { List<AgentToolSpec> operateToolSpecs) {
List<AgentToolSpec> merged = new ArrayList<>(); List<AgentToolSpec> merged = new ArrayList<>();
if (toolSpecs != null) { if (toolSpecs != null) {
merged.addAll(toolSpecs); merged.addAll(toolSpecs);
} }
if (knowledgeToolSpecs != null) {
merged.addAll(knowledgeToolSpecs);
}
if (mcpToolSpecs != null) { if (mcpToolSpecs != null) {
merged.addAll(mcpToolSpecs); merged.addAll(mcpToolSpecs);
} }
@@ -1251,6 +1232,46 @@ public class AgentScopeReActRuntime implements AgentRuntime {
return merged; return merged;
} }
/**
* 校验不同来源的运行时工具名称没有冲突。
*
* @param toolSpecs 普通工具声明
* @param knowledgeToolSpecs 知识库工具声明
* @param mcpToolSpecs MCP 工具声明
* @param operateToolSpecs 操作工具声明
* @throws AgentRuntimeException 工具名称重复时抛出
*/
private void validateRuntimeToolConflicts(List<AgentToolSpec> toolSpecs,
List<AgentToolSpec> knowledgeToolSpecs,
List<AgentToolSpec> mcpToolSpecs,
List<AgentToolSpec> operateToolSpecs) {
Set<String> names = new LinkedHashSet<>();
for (List<AgentToolSpec> specs : List.of(
safeToolSpecs(toolSpecs),
safeToolSpecs(knowledgeToolSpecs),
safeToolSpecs(mcpToolSpecs),
safeToolSpecs(operateToolSpecs))) {
for (AgentToolSpec spec : specs) {
if (spec == null || spec.getName() == null || spec.getName().isBlank()) {
continue;
}
if (!names.add(spec.getName())) {
throw new AgentRuntimeException("Agent runtime tool name conflict: " + spec.getName());
}
}
}
}
/**
* 将可空工具列表转换为空安全列表。
*
* @param toolSpecs 工具声明
* @return 非空工具声明列表
*/
private List<AgentToolSpec> safeToolSpecs(List<AgentToolSpec> toolSpecs) {
return toolSpecs == null ? List.of() : toolSpecs;
}
private void closeMcpClients() { private void closeMcpClients() {
for (McpClientWrapper client : mcpClients) { for (McpClientWrapper client : mcpClients) {
if (client == null) { if (client == null) {
@@ -1264,28 +1285,6 @@ public class AgentScopeReActRuntime implements AgentRuntime {
mcpClients.clear(); mcpClients.clear();
} }
/**
* 构建聚合知识库的默认检索配置。
*
* @param definition 智能体定义
* @return 检索配置
*/
private RetrieveConfig defaultRetrieveConfig(AgentDefinition definition) {
int limit = definition.getKnowledgeSpecs().stream()
.mapToInt(AgentKnowledgeSpec::getLimit)
.filter(value -> value > 0)
.sum();
double scoreThreshold = definition.getKnowledgeSpecs().stream()
.mapToDouble(AgentKnowledgeSpec::getScoreThreshold)
.filter(value -> value > 0D)
.min()
.orElse(0D);
return RetrieveConfig.builder()
.limit(limit <= 0 ? 5 : limit)
.scoreThreshold(scoreThreshold)
.build();
}
public AgentInitRequest getInitRequest() { public AgentInitRequest getInitRequest() {
return initRequest; return initRequest;
} }
@@ -1300,6 +1299,7 @@ public class AgentScopeReActRuntime implements AgentRuntime {
} }
private record AgentScopeToolkitBuildResult(Map<String, List<AgentTool>> skillTools, private record AgentScopeToolkitBuildResult(Map<String, List<AgentTool>> skillTools,
List<AgentToolSpec> knowledgeToolSpecs,
List<AgentToolSpec> mcpToolSpecs, List<AgentToolSpec> mcpToolSpecs,
List<AgentToolSpec> operateToolSpecs, List<AgentToolSpec> operateToolSpecs,
List<McpSkillRegistration> skillMcpRegistrations) { List<McpSkillRegistration> skillMcpRegistrations) {

View File

@@ -170,7 +170,8 @@ public class AgentScopeToolAdapter {
throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName()); throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName());
} }
return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder, return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder,
skillContext, skillBinding, emitNormalToolResult, true, true); skillContext, skillBinding, emitNormalToolResult, true, true,
resolveInvocationClassLoader(invoker));
} }
/** /**
@@ -203,7 +204,8 @@ public class AgentScopeToolAdapter {
throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName()); throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName());
} }
return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder, return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder,
skillContext, skillBinding, emitNormalToolResult, emitSkillStep, true); skillContext, skillBinding, emitNormalToolResult, emitSkillStep, true,
resolveInvocationClassLoader(invoker));
} }
/** /**
@@ -238,7 +240,24 @@ public class AgentScopeToolAdapter {
throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName()); throw new AgentRuntimeException("Agent tool invoker is required: " + toolSpec.getName());
} }
return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder, return new RuntimeAgentTool(toolSpec, invoker, request, approvalCoordinator, turnContextHolder,
skillContext, skillBinding, emitNormalToolResult, emitSkillStep, handleApprovalInTool); skillContext, skillBinding, emitNormalToolResult, emitSkillStep, handleApprovalInTool,
resolveInvocationClassLoader(invoker));
}
/**
* 解析工具执行时应使用的应用类加载器。
*
* <p>AgentScope 可能在 Reactor 工作线程执行工具。Spring Boot 可执行包中的业务类依赖
* 注册工具时的应用类加载器,不能依赖工作线程可能继承到的系统类加载器。</p>
*
* @param invoker 工具调用器
* @return 工具调用器所属类加载器;无法取得时回退到当前线程上下文类加载器
*/
private ClassLoader resolveInvocationClassLoader(AgentToolInvoker invoker) {
ClassLoader invokerClassLoader = invoker.getClass().getClassLoader();
return invokerClassLoader == null
? Thread.currentThread().getContextClassLoader()
: invokerClassLoader;
} }
private AgentRuntimeTurnContextHolder fixedHolder(AgentRuntimeExecutionContext request, private AgentRuntimeTurnContextHolder fixedHolder(AgentRuntimeExecutionContext request,
@@ -258,7 +277,8 @@ public class AgentScopeToolAdapter {
AgentSkillBinding skillBinding, AgentSkillBinding skillBinding,
boolean emitNormalToolResult, boolean emitNormalToolResult,
boolean emitSkillStep, boolean emitSkillStep,
boolean handleApprovalInTool) implements AgentTool { boolean handleApprovalInTool,
ClassLoader invocationClassLoader) implements AgentTool {
/** /**
* 获取工具名称。 * 获取工具名称。
@@ -402,6 +422,14 @@ public class AgentScopeToolAdapter {
* @return 工具结果块 * @return 工具结果块
*/ */
private ToolResultBlock invokeTool(ToolCallParam param, Map<String, Object> input) { private ToolResultBlock invokeTool(ToolCallParam param, Map<String, Object> input) {
Thread currentThread = Thread.currentThread();
ClassLoader originalClassLoader = currentThread.getContextClassLoader();
boolean switchClassLoader = invocationClassLoader != null
&& invocationClassLoader != originalClassLoader;
if (switchClassLoader) {
currentThread.setContextClassLoader(invocationClassLoader);
}
try {
AgentToolContext context = buildContext(param); AgentToolContext context = buildContext(param);
AgentToolResult result = invoker.invoke(input, context); AgentToolResult result = invoker.invoke(input, context);
ToolResultBlock block = toToolResultBlock(param, result); ToolResultBlock block = toToolResultBlock(param, result);
@@ -411,6 +439,11 @@ public class AgentScopeToolAdapter {
emit(toolResultEvent(block)); emit(toolResultEvent(block));
} }
return block; return block;
} finally {
if (switchClassLoader) {
currentThread.setContextClassLoader(originalClassLoader);
}
}
} }
/** /**

View File

@@ -107,9 +107,9 @@ public class AgentRuntimeTurnContext {
merged.setUserMessage(executionContext.getUserMessage()); merged.setUserMessage(executionContext.getUserMessage());
merged.setMemorySnapshot(executionContext.getMemorySnapshot()); merged.setMemorySnapshot(executionContext.getMemorySnapshot());
merged.setToolInvokers(fallback == null ? executionContext.getToolInvokers() : fallback.getToolInvokers()); merged.setToolInvokers(fallback == null ? executionContext.getToolInvokers() : fallback.getToolInvokers());
merged.setKnowledgeRetrievers(fallback == null merged.setKnowledgeRegistrations(fallback == null
? executionContext.getKnowledgeRetrievers() ? executionContext.getKnowledgeRegistrations()
: fallback.getKnowledgeRetrievers()); : fallback.getKnowledgeRegistrations());
merged.setSessionStore(fallback == null ? executionContext.getSessionStore() : fallback.getSessionStore()); merged.setSessionStore(fallback == null ? executionContext.getSessionStore() : fallback.getSessionStore());
merged.setConversationRecorder(fallback == null merged.setConversationRecorder(fallback == null
? executionContext.getConversationRecorder() ? executionContext.getConversationRecorder()

View File

@@ -5,10 +5,12 @@ import com.easyagents.agent.runtime.event.AgentRuntimeEventBridge;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType; import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agent.runtime.event.AgentRuntimeObserver; import com.easyagents.agent.runtime.event.AgentRuntimeObserver;
import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext; import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext;
import com.easyagents.agent.runtime.tool.AgentToolCategory;
import com.easyagents.agent.runtime.tool.AgentToolSpec; import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.hook.HookEvent; import io.agentscope.core.hook.HookEvent;
import io.agentscope.core.hook.PostActingEvent; import io.agentscope.core.hook.PostActingEvent;
import io.agentscope.core.hook.PreActingEvent; import io.agentscope.core.hook.PreActingEvent;
import io.agentscope.core.message.TextBlock;
import io.agentscope.core.message.ToolResultBlock; import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock; import io.agentscope.core.message.ToolUseBlock;
import reactor.core.publisher.Mono; import reactor.core.publisher.Mono;
@@ -130,12 +132,21 @@ public class ToolExecutionObserver implements AgentRuntimeObserver {
private void enrichToolPayload(AgentRuntimeEvent runtimeEvent, String toolName) { private void enrichToolPayload(AgentRuntimeEvent runtimeEvent, String toolName) {
AgentToolSpec toolSpec = toolSpecs.get(toolName); AgentToolSpec toolSpec = toolSpecs.get(toolName);
if (toolSpec == null || toolSpec.getMetadata() == null || toolSpec.getMetadata().isEmpty()) { if (toolSpec == null) {
return;
}
runtimeEvent.getPayload().put("toolCategory", toolSpec.getCategory().name());
if (toolSpec.getMetadata() == null || toolSpec.getMetadata().isEmpty()) {
return; return;
} }
Map<String, Object> metadata = toolSpec.getMetadata(); Map<String, Object> metadata = toolSpec.getMetadata();
putIfPresent(runtimeEvent.getPayload(), metadata, "toolDisplayName"); putIfPresent(runtimeEvent.getPayload(), metadata, "toolDisplayName");
putIfPresent(runtimeEvent.getPayload(), metadata, "skillId"); putIfPresent(runtimeEvent.getPayload(), metadata, "skillId");
if (toolSpec.getCategory() == AgentToolCategory.KNOWLEDGE) {
putIfPresent(runtimeEvent.getPayload(), metadata, "knowledgeId");
putIfPresent(runtimeEvent.getPayload(), metadata, "knowledgeName");
putIfPresent(runtimeEvent.getPayload(), metadata, "knowledgeRuntimeName");
}
} }
private void putIfPresent(Map<String, Object> payload, Map<String, Object> metadata, String key) { private void putIfPresent(Map<String, Object> payload, Map<String, Object> metadata, String key) {
@@ -149,7 +160,15 @@ public class ToolExecutionObserver implements AgentRuntimeObserver {
return false; return false;
} }
Object success = result.getMetadata() == null ? null : result.getMetadata().get("success"); Object success = result.getMetadata() == null ? null : result.getMetadata().get("success");
return !(success instanceof Boolean) || Boolean.TRUE.equals(success); if (success instanceof Boolean) {
return Boolean.TRUE.equals(success);
}
// AgentScope 1.x 将工具异常转换为不带 success metadata 的 "Error: ..." 文本结果。
return result.getOutput().stream()
.filter(TextBlock.class::isInstance)
.map(TextBlock.class::cast)
.map(TextBlock::getText)
.noneMatch(text -> text != null && text.startsWith("Error: "));
} }
private boolean isSkillTool(String toolName) { private boolean isSkillTool(String toolName) {

View File

@@ -1,10 +0,0 @@
package com.easyagents.agent.runtime.knowledge;
/**
* 知识库检索策略。
*/
public enum AgentKnowledgePolicy {
AGENTIC,
GENERIC,
DISABLED
}

View File

@@ -0,0 +1,43 @@
package com.easyagents.agent.runtime.knowledge;
import java.util.Objects;
/**
* 知识库声明与检索器的运行时绑定。
*/
public final class AgentKnowledgeRegistration {
private final AgentKnowledgeSpec knowledgeSpec;
private final AgentKnowledgeRetriever retriever;
/**
* 创建知识库运行时绑定。
*
* @param knowledgeSpec 知识库声明
* @param retriever 知识库检索器
* @throws NullPointerException 声明或检索器为空时抛出
*/
public AgentKnowledgeRegistration(AgentKnowledgeSpec knowledgeSpec,
AgentKnowledgeRetriever retriever) {
this.knowledgeSpec = Objects.requireNonNull(knowledgeSpec, "knowledgeSpec");
this.retriever = Objects.requireNonNull(retriever, "retriever");
}
/**
* 获取知识库声明。
*
* @return 知识库声明
*/
public AgentKnowledgeSpec getKnowledgeSpec() {
return knowledgeSpec;
}
/**
* 获取知识库检索器。
*
* @return 知识库检索器
*/
public AgentKnowledgeRetriever getRetriever() {
return retriever;
}
}

View File

@@ -9,9 +9,9 @@ import java.util.Map;
public class AgentKnowledgeSpec { public class AgentKnowledgeSpec {
private String knowledgeId; private String knowledgeId;
private String runtimeName;
private String name; private String name;
private String description; private String description;
private AgentKnowledgePolicy retrievalMode = AgentKnowledgePolicy.AGENTIC;
private int limit = 5; private int limit = 5;
private double scoreThreshold = 0D; private double scoreThreshold = 0D;
private Map<String, Object> metadata = new LinkedHashMap<>(); private Map<String, Object> metadata = new LinkedHashMap<>();
@@ -34,6 +34,24 @@ public class AgentKnowledgeSpec {
this.knowledgeId = knowledgeId; this.knowledgeId = knowledgeId;
} }
/**
* 获取调用方提供的英文运行名。
*
* @return 英文运行名
*/
public String getRuntimeName() {
return runtimeName;
}
/**
* 设置调用方提供的英文运行名。
*
* @param runtimeName 英文运行名
*/
public void setRuntimeName(String runtimeName) {
this.runtimeName = runtimeName;
}
/** /**
* 获取知识库名称。 * 获取知识库名称。
* *
@@ -70,24 +88,6 @@ public class AgentKnowledgeSpec {
this.description = description; this.description = description;
} }
/**
* 获取检索模式。
*
* @return 检索模式
*/
public AgentKnowledgePolicy getRetrievalMode() {
return retrievalMode;
}
/**
* 设置检索模式。
*
* @param retrievalMode 检索模式
*/
public void setRetrievalMode(AgentKnowledgePolicy retrievalMode) {
this.retrievalMode = retrievalMode == null ? AgentKnowledgePolicy.AGENTIC : retrievalMode;
}
/** /**
* 获取限制数量。 * 获取限制数量。
* *

View File

@@ -0,0 +1,79 @@
package com.easyagents.agent.runtime.knowledge;
import com.easyagents.agent.runtime.AgentRuntimeException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.HexFormat;
import java.util.regex.Pattern;
/**
* Agent 知识库工具名称规范。
*/
public final class AgentKnowledgeToolNames {
/**
* 知识库工具名称前缀。
*/
public static final String PREFIX = "retrieve_knowledge_";
/**
* OpenAI-compatible Function Call 的通用名称长度上限。
*/
public static final int MAX_TOOL_NAME_LENGTH = 64;
private static final int HASH_LENGTH = 8;
private static final Pattern SAFE_RUNTIME_NAME = Pattern.compile("^[A-Za-z0-9_-]+$");
/**
* 禁止实例化工具类。
*/
private AgentKnowledgeToolNames() {
}
/**
* 根据调用方运行名生成稳定的知识库工具名。
*
* <p>超长名称保留可读前缀并追加稳定短哈希,避免不同模型服务对 Function Call
* 名称长度限制不一致。</p>
*
* @param runtimeName 调用方提供的英文运行名
* @return 合法且稳定的工具名
* @throws AgentRuntimeException 运行名为空或包含非法字符时抛出
*/
public static String build(String runtimeName) {
String normalized = runtimeName == null ? "" : runtimeName.trim();
if (normalized.isEmpty()) {
throw new AgentRuntimeException("Knowledge runtime name is required.");
}
if (!SAFE_RUNTIME_NAME.matcher(normalized).matches()) {
throw new AgentRuntimeException(
"Knowledge runtime name must contain only letters, numbers, underscores, or hyphens: "
+ normalized);
}
String toolName = PREFIX + normalized;
if (toolName.length() <= MAX_TOOL_NAME_LENGTH) {
return toolName;
}
String hash = shortHash(normalized);
int readableLength = MAX_TOOL_NAME_LENGTH - PREFIX.length() - HASH_LENGTH - 1;
return PREFIX + normalized.substring(0, readableLength) + "_" + hash;
}
/**
* 计算用于超长工具名消歧的稳定短哈希。
*
* @param value 原始运行名
* @return 八位十六进制哈希
*/
private static String shortHash(String value) {
try {
byte[] digest = MessageDigest.getInstance("SHA-256")
.digest(value.getBytes(StandardCharsets.UTF_8));
return HexFormat.of().formatHex(digest).substring(0, HASH_LENGTH);
} catch (NoSuchAlgorithmException exception) {
throw new IllegalStateException("SHA-256 is unavailable.", exception);
}
}
}

View File

@@ -34,9 +34,16 @@ public class HeuristicKnowledgeCitationMatcher implements AgentKnowledgeCitation
if (normalizedAnswer.length() < MIN_NORMALIZED_ANSWER_LENGTH) { if (normalizedAnswer.length() < MIN_NORMALIZED_ANSWER_LENGTH) {
return List.of(); return List.of();
} }
List<String> normalizedSegments = normalizeSegments(answerText);
List<ScoredKnowledgeReference> scoredReferences = new ArrayList<>(); List<ScoredKnowledgeReference> scoredReferences = new ArrayList<>();
for (AgentKnowledgeReference candidate : candidates) { for (AgentKnowledgeReference candidate : candidates) {
double supportScore = supportScore(normalizedAnswer, normalize(candidate == null ? null : candidate.getChunkContent())); String normalizedContent = normalize(candidate == null ? null : candidate.getChunkContent());
double supportScore = supportScore(normalizedAnswer, normalizedContent);
// 长篇汇总回答中,每条证据通常只支撑一个段落或条目。继续只用整篇答案作分母,
// 会随着主题增多把有效引用的重合度稀释到阈值以下。
for (String normalizedSegment : normalizedSegments) {
supportScore = Math.max(supportScore, supportScore(normalizedSegment, normalizedContent));
}
if (supportScore >= MIN_SUPPORT_SCORE) { if (supportScore >= MIN_SUPPORT_SCORE) {
scoredReferences.add(new ScoredKnowledgeReference(candidate, supportScore)); scoredReferences.add(new ScoredKnowledgeReference(candidate, supportScore));
} }
@@ -48,6 +55,22 @@ public class HeuristicKnowledgeCitationMatcher implements AgentKnowledgeCitation
.toList(); .toList();
} }
/**
* 将答案切分为可独立核验的段落或句子并完成归一化。
*
* @param answerText 最终答案文本
* @return 非空的归一化答案片段
*/
private List<String> normalizeSegments(String answerText) {
if (answerText == null || answerText.isBlank()) {
return List.of();
}
return Arrays.stream(answerText.split("[\\r\\n。!?;]+"))
.map(this::normalize)
.filter(segment -> segment.length() >= MIN_NORMALIZED_ANSWER_LENGTH)
.toList();
}
/** /**
* 计算答案与候选片段之间的文本支撑分。 * 计算答案与候选片段之间的文本支撑分。
* *

View File

@@ -0,0 +1,105 @@
package com.easyagents.agent.runtime.prompt;
import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import java.util.ArrayList;
import java.util.List;
/**
* 组合智能体运行时使用的系统提示词。
*/
public final class SystemPromptComposer {
private static final String KNOWLEDGE_TOOL_PROTOCOL = """
Knowledge tool protocol:
- Knowledge tools provide grounded information for the scopes described in their tool descriptions.
- Before answering a request involving facts, policies, procedures, entitlements, prices, product or service details, or other information within a knowledge tool's scope, call the most relevant knowledge tool first.
- Do not rely only on general model knowledge for claims that fall within an available knowledge tool's scope.
- If it is uncertain whether the request falls within a knowledge tool's scope, prefer making one retrieval call.
- Build a concise, standalone query from the current request and only the necessary conversation context. Resolve pronouns and omitted subjects, and preserve relevant names, time, location, product, membership level, constraints, and user intent. Do not copy the entire conversation.
- If the returned results are empty or do not address the request, reformulate the query and retry once when a materially different query is possible. Do not repeat the same query or retrieve indefinitely.
- Base knowledge-backed claims only on information actually returned by the tools. Never imply that a knowledge base contains information that was not returned.
- If retrieval remains insufficient, follow the Agent's configured system prompt for the response strategy.
- Skip retrieval for greetings, casual conversation, pure writing or translation, and tasks that clearly do not depend on knowledge-base facts.
""".strip();
private static final String ASYNC_TOOL_PROTOCOL = """
Async tool protocol:
- Async tools may expose submit, observe, result, cancel, and list sub-tools. Treat these sub-tools as one user-facing tool.
- Do not ask the user to choose submit, observe, result, cancel, or list. These are internal execution phases.
- For a normal user request to use an async tool, call its submit sub-tool first with the user-provided arguments by default.
- After submit returns task_id, immediately call observe with that task_id to check progress.
- If the task is completed and result is available, use the returned result to answer the user.
- If the task is still running after observation, tell the user that the task is running and keep task_id/next_action for later tool calls.
- Use result, list, or cancel directly only when the user explicitly asks to get a known task result, list tasks, or cancel a task.
""".strip();
/**
* 阻止工具类被实例化。
*/
private SystemPromptComposer() {
}
/**
* 按固定顺序组合用户提示词与运行时工具协议。
*
* @param definition 智能体定义
* @return 最终系统提示词;没有任何提示词时返回原始空值
*/
public static String compose(AgentDefinition definition) {
if (definition == null) {
return null;
}
boolean knowledgeProtocolEnabled = hasEnabledKnowledgeTool(definition);
boolean asyncProtocolEnabled = hasAsyncTool(definition);
String userPrompt = definition.getSystemPrompt();
if (!knowledgeProtocolEnabled && !asyncProtocolEnabled) {
return userPrompt;
}
List<String> promptSections = new ArrayList<>(3);
if (userPrompt != null && !userPrompt.isBlank()) {
promptSections.add(userPrompt.stripTrailing());
}
if (knowledgeProtocolEnabled) {
promptSections.add(KNOWLEDGE_TOOL_PROTOCOL);
}
if (asyncProtocolEnabled) {
promptSections.add(ASYNC_TOOL_PROTOCOL);
}
return String.join("\n\n", promptSections);
}
/**
* 判断当前定义是否会注册可调用的知识库工具。
*
* @param definition 智能体定义
* @return 知识库工具可用时返回 true
*/
private static boolean hasEnabledKnowledgeTool(AgentDefinition definition) {
return definition.getExecutionOptions() != null
&& definition.getExecutionOptions().isToolCallingEnabled()
&& definition.getKnowledgeSpecs() != null
&& !definition.getKnowledgeSpecs().isEmpty();
}
/**
* 判断当前定义是否包含异步工具。
*
* @param definition 智能体定义
* @return 包含异步工具时返回 true
*/
private static boolean hasAsyncTool(AgentDefinition definition) {
if (definition.getToolSpecs() == null) {
return false;
}
for (AgentToolSpec toolSpec : definition.getToolSpecs()) {
// AsyncToolSpecExpander 会为生成的全部子工具写入该运行时元数据。
if (toolSpec != null && Boolean.TRUE.equals(toolSpec.getMetadata().get("asyncTool"))) {
return true;
}
}
return false;
}
}

View File

@@ -0,0 +1,226 @@
package com.easyagents.agent.runtime.agentscope;
import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.event.AgentRuntimeEvent;
import com.easyagents.agent.runtime.event.AgentRuntimeEventBridge;
import com.easyagents.agent.runtime.event.AgentRuntimeEventType;
import com.easyagents.agent.runtime.event.AgentRuntimeTurnContext;
import com.easyagents.agent.runtime.event.AgentRuntimeTurnContextHolder;
import com.easyagents.agent.runtime.hitl.AgentToolApprovalCoordinator;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeDocument;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalResult;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeToolNames;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.tool.ToolCallParam;
import io.agentscope.core.tool.Toolkit;
import org.junit.Assert;
import org.junit.Test;
import reactor.core.publisher.Sinks;
import java.time.Duration;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;
/**
* {@link AgentScopeKnowledgeAdapter} 回归测试。
*/
public class AgentScopeKnowledgeAdapterTest {
/**
* 验证每个知识库生成独立工具,工具名、描述和 Schema 向模型暴露完整信息。
*/
@Test
public void shouldCreateOneToolPerKnowledge() {
AgentKnowledgeSpec first = knowledgeSpec("knowledge-1", "homeinn_faq", 5);
first.setName("如家 FAQ");
first.setDescription("如家酒店入住、退房和会员服务常见问题");
AgentKnowledgeSpec second = knowledgeSpec("knowledge-2", "hotel_policy", 3);
AgentRuntimeExecutionContext context = executionContext(List.of(first, second));
context.setKnowledgeRegistrations(List.of(
registration(first, 1, 0.9D),
registration(second, 1, 0.8D)));
List<AgentToolSpec> toolSpecs = new AgentScopeKnowledgeAdapter().createToolSpecs(context);
Assert.assertEquals(2, toolSpecs.size());
Assert.assertEquals("retrieve_knowledge_homeinn_faq", toolSpecs.get(0).getName());
Assert.assertTrue(toolSpecs.get(0).getDescription().contains("如家 FAQ"));
Assert.assertTrue(toolSpecs.get(0).getDescription().contains("入住、退房"));
Assert.assertEquals(List.of("query"), toolSpecs.get(0).getParametersSchema().get("required"));
Assert.assertFalse(String.valueOf(toolSpecs.get(0).getParametersSchema()).contains("limit"));
}
/**
* 验证选中一个知识库工具时只调用对应 Retriever并使用绑定 limit 和阈值。
*/
@Test
public void shouldDispatchOnlyToSelectedKnowledgeRetriever() {
AgentKnowledgeSpec first = knowledgeSpec("knowledge-1", "homeinn_faq", 2);
first.setScoreThreshold(0.5D);
AgentKnowledgeSpec second = knowledgeSpec("knowledge-2", "hotel_policy", 2);
AgentRuntimeExecutionContext context = executionContext(List.of(first, second));
AtomicInteger firstCalls = new AtomicInteger();
AtomicInteger secondCalls = new AtomicInteger();
context.setKnowledgeRegistrations(List.of(
new AgentKnowledgeRegistration(first, request -> {
firstCalls.incrementAndGet();
Assert.assertEquals("如家几点退房", request.getQuery());
Assert.assertEquals(2, request.getLimit());
Assert.assertEquals(0.5D, request.getScoreThreshold(), 0.0001D);
return AgentKnowledgeRetrievalResult.of(List.of(
document("knowledge-1", 0.9D),
document("knowledge-1-low", 0.2D)));
}),
new AgentKnowledgeRegistration(second, request -> {
secondCalls.incrementAndGet();
return AgentKnowledgeRetrievalResult.of(List.of(document("knowledge-2", 0.8D)));
})));
RegisteredKnowledgeTools registered = register(context);
ToolResultBlock result = registered.toolkit().getTool("retrieve_knowledge_homeinn_faq")
.callAsync(toolCall("retrieve_knowledge_homeinn_faq", "如家几点退房"))
.block();
Assert.assertNotNull(result);
Assert.assertEquals(1, firstCalls.get());
Assert.assertEquals(0, secondCalls.get());
Assert.assertEquals(1, result.getMetadata().get("documentCount"));
}
/**
* 验证最终知识库事件与阈值过滤、排序和截断后的模型可见文档一致。
*/
@Test
public void retrievalEventShouldMatchFinalToolDocuments() {
AgentKnowledgeSpec spec = knowledgeSpec("knowledge-1", "homeinn_faq", 2);
spec.setScoreThreshold(0.5D);
AgentRuntimeExecutionContext context = executionContext(List.of(spec));
context.setKnowledgeRegistrations(List.of(new AgentKnowledgeRegistration(spec, request ->
AgentKnowledgeRetrievalResult.of(List.of(
document("third", 0.7D),
document("first", 0.9D),
document("filtered", 0.2D),
document("second", 0.8D))))));
RegisteredKnowledgeTools registered = register(context);
ToolResultBlock result = registered.toolkit().getTool("retrieve_knowledge_homeinn_faq")
.callAsync(toolCall("retrieve_knowledge_homeinn_faq", "query"))
.block();
List<AgentRuntimeEvent> events = registered.eventSink().asFlux()
.filter(event -> event.getEventType() == AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL)
.take(1)
.collectList()
.block(Duration.ofSeconds(1));
Assert.assertNotNull(result);
Assert.assertNotNull(events);
Assert.assertEquals(1, events.size());
AgentRuntimeEvent event = events.get(0);
Assert.assertEquals(2, event.getPayload().get("documentCount"));
@SuppressWarnings("unchecked")
List<Map<String, Object>> eventDocuments =
(List<Map<String, Object>>) event.getPayload().get("documents");
@SuppressWarnings("unchecked")
List<Map<String, Object>> resultDocuments =
(List<Map<String, Object>>) result.getMetadata().get("documents");
Assert.assertEquals(resultDocuments, eventDocuments);
Assert.assertEquals("first", eventDocuments.get(0).get("documentId"));
Assert.assertEquals("second", eventDocuments.get(1).get("documentId"));
}
/**
* 验证超长运行名会稳定压缩到 Function Call 长度上限。
*/
@Test
public void longRuntimeNameShouldProduceStableBoundedToolName() {
String runtimeName = "knowledge_" + "a".repeat(80);
String first = AgentKnowledgeToolNames.build(runtimeName);
String second = AgentKnowledgeToolNames.build(runtimeName);
Assert.assertEquals(first, second);
Assert.assertEquals(AgentKnowledgeToolNames.MAX_TOOL_NAME_LENGTH, first.length());
Assert.assertTrue(first.startsWith(AgentKnowledgeToolNames.PREFIX));
}
private RegisteredKnowledgeTools register(AgentRuntimeExecutionContext context) {
AgentScopeKnowledgeAdapter adapter = new AgentScopeKnowledgeAdapter();
List<AgentToolSpec> toolSpecs = adapter.createToolSpecs(context);
Toolkit toolkit = new Toolkit();
Sinks.Many<AgentRuntimeEvent> eventSink = Sinks.many().replay().all();
AgentRuntimeTurnContextHolder holder = new AgentRuntimeTurnContextHolder();
AgentRuntimeEventBridge eventBridge = new AgentRuntimeEventBridge(context, holder);
holder.set(new AgentRuntimeTurnContext(context, eventSink, eventBridge));
adapter.registerTools(context, toolSpecs, toolkit, new AgentScopeToolAdapter(),
AgentToolApprovalCoordinator.disabled(), holder);
return new RegisteredKnowledgeTools(toolkit, eventSink);
}
private ToolCallParam toolCall(String toolName, String query) {
ToolUseBlock toolUseBlock = ToolUseBlock.builder()
.id("call-1")
.name(toolName)
.input(Map.of("query", query))
.build();
return ToolCallParam.builder()
.toolUseBlock(toolUseBlock)
.input(toolUseBlock.getInput())
.build();
}
private AgentRuntimeExecutionContext executionContext(List<AgentKnowledgeSpec> specs) {
AgentDefinition definition = new AgentDefinition();
definition.setAgentId("agent-1");
definition.setKnowledgeSpecs(specs);
AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext();
context.setAgentDefinition(definition);
context.setTraceId("trace-1");
context.setSessionId("session-1");
return context;
}
private AgentKnowledgeSpec knowledgeSpec(String knowledgeId, String runtimeName, int limit) {
AgentKnowledgeSpec spec = new AgentKnowledgeSpec();
spec.setKnowledgeId(knowledgeId);
spec.setRuntimeName(runtimeName);
spec.setName(knowledgeId);
spec.setLimit(limit);
return spec;
}
private AgentKnowledgeRegistration registration(AgentKnowledgeSpec spec,
int documentCount,
double firstScore) {
return new AgentKnowledgeRegistration(spec, request ->
AgentKnowledgeRetrievalResult.of(documents(spec.getKnowledgeId(), documentCount, firstScore)));
}
private List<AgentKnowledgeDocument> documents(String knowledgeId, int count, double firstScore) {
List<AgentKnowledgeDocument> documents = new ArrayList<>();
for (int index = 0; index < count; index++) {
documents.add(document(knowledgeId + "-doc-" + index, firstScore - (index * 0.01D)));
}
return documents;
}
private AgentKnowledgeDocument document(String documentId, double score) {
AgentKnowledgeDocument document = new AgentKnowledgeDocument();
document.setDocumentId(documentId);
document.setDocumentName("FAQ");
document.setChunkId(documentId + "-chunk");
document.setContent("content-" + documentId);
document.setScore(score);
return document;
}
private record RegisteredKnowledgeTools(Toolkit toolkit,
Sinks.Many<AgentRuntimeEvent> eventSink) {
}
}

View File

@@ -10,6 +10,7 @@ import com.easyagents.agent.runtime.event.observer.SkillExecutionObserver;
import com.easyagents.agent.runtime.event.observer.ToolExecutionObserver; import com.easyagents.agent.runtime.event.observer.ToolExecutionObserver;
import com.easyagents.agent.runtime.hitl.AgentResumeToken; import com.easyagents.agent.runtime.hitl.AgentResumeToken;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeDocument; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeDocument;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRegistration;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalResult; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeRetrievalResult;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec; import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.memory.AgentMemoryCompressionParameter; import com.easyagents.agent.runtime.memory.AgentMemoryCompressionParameter;
@@ -23,6 +24,7 @@ import com.easyagents.agent.runtime.persistence.session.memory.InMemoryAgentSess
import com.easyagents.agent.runtime.skill.AgentSkillBoxSpec; import com.easyagents.agent.runtime.skill.AgentSkillBoxSpec;
import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext; import com.easyagents.agent.runtime.skill.AgentSkillRuntimeContext;
import com.easyagents.agent.runtime.skill.AgentSkillSpec; import com.easyagents.agent.runtime.skill.AgentSkillSpec;
import com.easyagents.agent.runtime.tool.AgentToolCategory;
import com.easyagents.agent.runtime.tool.AgentToolResult; import com.easyagents.agent.runtime.tool.AgentToolResult;
import com.easyagents.agent.runtime.tool.AgentToolSpec; import com.easyagents.agent.runtime.tool.AgentToolSpec;
import com.easyagents.agent.runtime.tool.operate.AgentOperateToolAdapter; import com.easyagents.agent.runtime.tool.operate.AgentOperateToolAdapter;
@@ -153,6 +155,27 @@ public class AgentScopeStatefulRuntimeTest {
Assert.assertTrue(runtime.getAgent().getSysPrompt().contains("immediately call observe")); Assert.assertTrue(runtime.getAgent().getSysPrompt().contains("immediately call observe"));
} }
@Test
public void shouldAppendKnowledgeToolProtocolPromptWhenKnowledgeToolsExist() {
AgentInitRequest request = initRequest();
AgentKnowledgeSpec knowledgeSpec = new AgentKnowledgeSpec();
knowledgeSpec.setKnowledgeId("knowledge-faq");
knowledgeSpec.setRuntimeName("faq");
request.getAgentDefinition().setKnowledgeSpecs(List.of(knowledgeSpec));
request.setKnowledgeRegistrations(List.of(new AgentKnowledgeRegistration(
knowledgeSpec, retrievalRequest -> AgentKnowledgeRetrievalResult.of(List.of()))));
AgentScopeReActRuntime runtime = fakeRuntime();
runtime.init(request);
Assert.assertTrue(runtime.getAgent().getSysPrompt()
.startsWith("system\n\nKnowledge tool protocol:"));
Assert.assertTrue(runtime.getAgent().getSysPrompt()
.contains("call the most relevant knowledge tool first"));
Assert.assertTrue(runtime.getAgent().getSysPrompt()
.contains("reformulate the query and retry once"));
}
@Test @Test
public void shouldEmitSideEventWithRuntimeIdentityFromBridge() throws Exception { public void shouldEmitSideEventWithRuntimeIdentityFromBridge() throws Exception {
AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext(); AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext();
@@ -235,6 +258,26 @@ public class AgentScopeStatefulRuntimeTest {
Assert.assertTrue(interceptors.stream().anyMatch(ToolHitlInterceptor.class::isInstance)); Assert.assertTrue(interceptors.stream().anyMatch(ToolHitlInterceptor.class::isInstance));
} }
/**
* 验证调用方提供的历史快照会在 Agent 首次构建时进入模型记忆。
*/
@Test
public void shouldAttachInitialConversationHistoryToAgentMemory() {
AgentInitRequest request = initRequest();
AgentMemorySnapshot snapshot = new AgentMemorySnapshot();
snapshot.addMessage(AgentMessage.text(AgentMessageRole.USER, "previous user question"));
snapshot.addMessage(AgentMessage.text(AgentMessageRole.ASSISTANT, "previous assistant answer"));
request.setMemorySnapshot(snapshot);
AgentScopeReActRuntime runtime = fakeRuntime();
runtime.init(request);
List<Msg> messages = runtime.getAgent().getMemory().getMessages();
Assert.assertEquals(2, messages.size());
Assert.assertEquals("previous user question", messages.get(0).getTextContent());
Assert.assertEquals("previous assistant answer", messages.get(1).getTextContent());
}
@Test @Test
public void shouldRegisterOperateToolsIntoToolkit() { public void shouldRegisterOperateToolsIntoToolkit() {
AgentInitRequest request = initRequest(); AgentInitRequest request = initRequest();
@@ -554,6 +597,82 @@ public class AgentScopeStatefulRuntimeTest {
Assert.assertFalse(events.toString().contains("sentinel-secret")); Assert.assertFalse(events.toString().contains("sentinel-secret"));
} }
/**
* 验证知识库工具复用统一工具生命周期,并携带可供调用方投影的知识库分类。
*
* @throws Exception 等待旁路事件失败时抛出
*/
@Test
public void shouldEmitKnowledgeToolLifecycleFromToolExecutionObserver() throws Exception {
AgentRuntimeExecutionContext context = executionContext();
Sinks.Many<AgentRuntimeEvent> sink = Sinks.many().replay().all();
AgentRuntimeEventBridge bridge = AgentRuntimeEventBridge.fixed(context, sink);
AgentToolSpec toolSpec = new AgentToolSpec();
toolSpec.setName("retrieve_knowledge_homeinn_faq");
toolSpec.setCategory(AgentToolCategory.KNOWLEDGE);
toolSpec.getMetadata().put("toolDisplayName", "如家 FAQ");
toolSpec.getMetadata().put("knowledgeId", "knowledge-1");
toolSpec.getMetadata().put("knowledgeName", "如家 FAQ");
toolSpec.getMetadata().put("knowledgeRuntimeName", "homeinn_faq");
ToolExecutionObserver observer = new ToolExecutionObserver(bridge, null, List.of(toolSpec));
ReActAgent agent = initializedAgent();
Toolkit toolkit = agent.getToolkit();
ToolUseBlock toolUse = ToolUseBlock.builder()
.id("knowledge-call-1")
.name("retrieve_knowledge_homeinn_faq")
.input(Map.of("query", "几点退房"))
.build();
ToolResultBlock toolResult = ToolResultBlock.of(
"knowledge-call-1",
"retrieve_knowledge_homeinn_faq",
TextBlock.builder().text("中午十二点退房").build(),
Map.of("success", true));
observer.observe(new PreActingEvent(agent, toolkit, toolUse)).block();
observer.observe(new PostActingEvent(agent, toolkit, toolUse, toolResult)).block();
List<AgentRuntimeEvent> events = sink.asFlux().take(2).collectList().toFuture()
.get(3, TimeUnit.SECONDS);
Assert.assertEquals(AgentRuntimeEventType.TOOL_CALL, events.get(0).getEventType());
Assert.assertEquals("KNOWLEDGE", events.get(0).getPayload().get("toolCategory"));
Assert.assertEquals("knowledge-1", events.get(0).getPayload().get("knowledgeId"));
Assert.assertEquals("RUNNING", events.get(0).getPayload().get("status"));
Assert.assertEquals(AgentRuntimeEventType.TOOL_RESULT, events.get(1).getEventType());
Assert.assertEquals("KNOWLEDGE", events.get(1).getPayload().get("toolCategory"));
Assert.assertEquals("SUCCESS", events.get(1).getPayload().get("status"));
}
/**
* 验证 AgentScope 的文本错误结果会投影为失败工具状态。
*
* @throws Exception 等待旁路事件失败时抛出
*/
@Test
public void shouldMarkAgentScopeKnowledgeToolErrorAsFailed() throws Exception {
AgentRuntimeExecutionContext context = executionContext();
Sinks.Many<AgentRuntimeEvent> sink = Sinks.many().replay().all();
AgentRuntimeEventBridge bridge = AgentRuntimeEventBridge.fixed(context, sink);
AgentToolSpec toolSpec = new AgentToolSpec();
toolSpec.setName("retrieve_knowledge_homeinn_faq");
toolSpec.setCategory(AgentToolCategory.KNOWLEDGE);
ToolExecutionObserver observer = new ToolExecutionObserver(bridge, null, List.of(toolSpec));
ReActAgent agent = initializedAgent();
ToolUseBlock toolUse = ToolUseBlock.builder()
.id("knowledge-call-error")
.name("retrieve_knowledge_homeinn_faq")
.input(Map.of("query", "异常查询"))
.build();
ToolResultBlock toolResult = ToolResultBlock.error("retriever unavailable")
.withIdAndName("knowledge-call-error", "retrieve_knowledge_homeinn_faq");
observer.observe(new PostActingEvent(agent, agent.getToolkit(), toolUse, toolResult)).block();
AgentRuntimeEvent event = sink.asFlux().next().toFuture().get(3, TimeUnit.SECONDS);
Assert.assertEquals(AgentRuntimeEventType.TOOL_RESULT, event.getEventType());
Assert.assertEquals("FAILED", event.getPayload().get("status"));
Assert.assertEquals(Boolean.FALSE, event.getPayload().get("success"));
}
@Test @Test
public void shouldEmitSkillLifecycleEventsFromSkillExecutionObserver() throws Exception { public void shouldEmitSkillLifecycleEventsFromSkillExecutionObserver() throws Exception {
AgentRuntimeExecutionContext context = executionContext(); AgentRuntimeExecutionContext context = executionContext();
@@ -838,6 +957,59 @@ public class AgentScopeStatefulRuntimeTest {
&& "partial answer".equals(message.getTextContent()))); && "partial answer".equals(message.getTextContent())));
} }
@Test
public void shouldPersistUserMessageWhenModelFails() {
InMemoryAgentSessionStore sessionStore = new InMemoryAgentSessionStore();
AgentInitRequest request = initRequest();
request.setSessionStore(sessionStore);
AgentScopeReActRuntime runtime = runtimeWithError(new IllegalStateException("model unavailable"));
runtime.init(request);
List<AgentRuntimeEvent> events = runtime.stream(
AgentMessage.text(AgentMessageRole.USER, "remember this request"))
.collectList()
.block();
Assert.assertTrue(events.stream().anyMatch(event -> event.getEventType() == AgentRuntimeEventType.FAILED));
Assert.assertTrue(sessionStore.exists("session-1"));
AgentScopeReActRuntime restoredRuntime = fakeRuntime();
restoredRuntime.init(request);
Assert.assertTrue(restoredRuntime.getAgent().getMemory().getMessages().stream()
.anyMatch(message -> message.getRole() == MsgRole.USER
&& "remember this request".equals(message.getTextContent())));
}
@Test
public void shouldPersistUserMessageWhenReasoningOnlyStreamIsCancelled() throws Exception {
InMemoryAgentSessionStore sessionStore = new InMemoryAgentSessionStore();
AgentInitRequest request = initRequest();
request.setSessionStore(sessionStore);
AgentScopeReActRuntime runtime = runtimeWithModel(List.of(
ChatResponse.builder()
.id("reasoning-only")
.content(List.of(ThinkingBlock.builder().thinking("still thinking").build()))
.build()), Duration.ofSeconds(5));
runtime.init(request);
CompletableFuture<AgentRuntimeEvent> firstReasoning = new CompletableFuture<>();
reactor.core.Disposable disposable = runtime.stream(
AgentMessage.text(AgentMessageRole.USER, "remember reasoning request"))
.subscribe(event -> {
if (event.getEventType() == AgentRuntimeEventType.REASONING_DELTA) {
firstReasoning.complete(event);
}
}, firstReasoning::completeExceptionally);
firstReasoning.get(3, TimeUnit.SECONDS);
disposable.dispose();
awaitCondition(() -> sessionStore.exists("session-1"));
AgentScopeReActRuntime restoredRuntime = fakeRuntime();
restoredRuntime.init(request);
Assert.assertTrue(restoredRuntime.getAgent().getMemory().getMessages().stream()
.anyMatch(message -> message.getRole() == MsgRole.USER
&& "remember reasoning request".equals(message.getTextContent())));
}
@Test @Test
public void shouldNotDuplicateNormalToolEventsFromMainStream() { public void shouldNotDuplicateNormalToolEventsFromMainStream() {
AgentInitRequest request = initRequest(); AgentInitRequest request = initRequest();
@@ -1577,9 +1749,12 @@ public class AgentScopeStatefulRuntimeTest {
AgentInitRequest request = initRequest(); AgentInitRequest request = initRequest();
AgentKnowledgeSpec knowledgeSpec = new AgentKnowledgeSpec(); AgentKnowledgeSpec knowledgeSpec = new AgentKnowledgeSpec();
knowledgeSpec.setKnowledgeId("knowledge-1"); knowledgeSpec.setKnowledgeId("knowledge-1");
knowledgeSpec.setRuntimeName("faq");
knowledgeSpec.setName("知识库"); knowledgeSpec.setName("知识库");
request.getAgentDefinition().setKnowledgeSpecs(List.of(knowledgeSpec)); request.getAgentDefinition().setKnowledgeSpecs(List.of(knowledgeSpec));
request.setKnowledgeRetrievers(Map.of("knowledge-1", retrievalRequest -> { AtomicBoolean knowledgeInvoked = new AtomicBoolean(false);
request.setKnowledgeRegistrations(List.of(new AgentKnowledgeRegistration(knowledgeSpec, retrievalRequest -> {
knowledgeInvoked.set(true);
AgentKnowledgeDocument document = new AgentKnowledgeDocument(); AgentKnowledgeDocument document = new AgentKnowledgeDocument();
document.setDocumentId("doc-1"); document.setDocumentId("doc-1");
document.setDocumentName("说明文档"); document.setDocumentName("说明文档");
@@ -1587,9 +1762,27 @@ public class AgentScopeStatefulRuntimeTest {
document.setContent("fake answer"); document.setContent("fake answer");
document.setScore(0.9D); document.setScore(0.9D);
return AgentKnowledgeRetrievalResult.of(List.of(document)); return AgentKnowledgeRetrievalResult.of(List.of(document));
})); })));
AgentScopeReActRuntime runtime = fakeRuntime(); AgentScopeReActRuntime runtime = runtimeWithModel(List.of(
ChatResponse.builder()
.id("knowledge-tool-call")
.content(List.of(ToolUseBlock.builder()
.id("call-search")
.name("retrieve_knowledge_faq")
.input(Map.of("query", "fake answer"))
.content("{\"query\":\"fake answer\"}")
.build()))
.finishReason("tool_calls")
.build(),
ChatResponse.builder()
.id("knowledge-final")
.content(List.of(TextBlock.builder().text("fake answer").build()))
.finishReason("stop")
.build()));
runtime.init(request); runtime.init(request);
Assert.assertNotNull(runtime.getAgent().getToolkit().getTool("retrieve_knowledge_faq"));
Assert.assertTrue(runtime.getAgent().getToolkit().getToolSchemas().stream()
.anyMatch(schema -> "retrieve_knowledge_faq".equals(schema.getName())));
List<AgentRuntimeEvent> events = runtime.stream(AgentMessage.text(AgentMessageRole.USER, "query knowledge")) List<AgentRuntimeEvent> events = runtime.stream(AgentMessage.text(AgentMessageRole.USER, "query knowledge"))
.collectList() .collectList()
@@ -1600,19 +1793,44 @@ public class AgentScopeStatefulRuntimeTest {
.findFirst() .findFirst()
.orElseThrow(); .orElseThrow();
Assert.assertNotNull(completed.getMessage()); Assert.assertNotNull(completed.getMessage());
if (completed.getMessage().getKnowledgeReferences().isEmpty()) { Assert.assertTrue("knowledgeInvoked=" + knowledgeInvoked.get() + ", events="
/* + events.stream().map(AgentRuntimeEvent::getEventType).toList(),
* 当前 runtime 将知识库注册为 AgentScope AGENTIC RAG模型需要主动调用 events.stream().anyMatch(event ->
* retrieve_knowledge 才会产生 KNOWLEDGE_RETRIEVAL 旁路事件。fake model
* 不会调用该工具时,不应强行猜引用。
*/
Assert.assertFalse(events.stream().anyMatch(event ->
event.getEventType() == AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL)); event.getEventType() == AgentRuntimeEventType.KNOWLEDGE_RETRIEVAL));
return; Assert.assertTrue(events.stream().anyMatch(event ->
} event.getEventType() == AgentRuntimeEventType.TOOL_CALL
&& "KNOWLEDGE".equals(event.getPayload().get("toolCategory"))));
Assert.assertTrue(events.stream().anyMatch(event ->
event.getEventType() == AgentRuntimeEventType.TOOL_RESULT
&& "KNOWLEDGE".equals(event.getPayload().get("toolCategory"))));
Assert.assertEquals("chunk-1", completed.getMessage().getKnowledgeReferences().get(0).getChunkId()); Assert.assertEquals("chunk-1", completed.getMessage().getKnowledgeReferences().get(0).getChunkId());
} }
/**
* 验证知识库生成工具与普通工具同名时拒绝初始化,避免 Toolkit 静默覆盖。
*/
@Test(expected = AgentRuntimeException.class)
public void shouldRejectKnowledgeToolNameConflictWithRegularTool() {
AgentInitRequest request = initRequest();
AgentKnowledgeSpec knowledgeSpec = new AgentKnowledgeSpec();
knowledgeSpec.setKnowledgeId("knowledge-1");
knowledgeSpec.setRuntimeName("faq");
knowledgeSpec.setName("知识库");
request.getAgentDefinition().setKnowledgeSpecs(List.of(knowledgeSpec));
request.setKnowledgeRegistrations(List.of(new AgentKnowledgeRegistration(
knowledgeSpec,
retrievalRequest -> AgentKnowledgeRetrievalResult.of(List.of()))));
AgentToolSpec regularTool = new AgentToolSpec();
regularTool.setName("retrieve_knowledge_faq");
regularTool.setDescription("conflicting tool");
request.getAgentDefinition().setToolSpecs(List.of(regularTool));
request.setToolInvokers(Map.of(
regularTool.getName(),
(arguments, context) -> AgentToolResult.success("done")));
fakeRuntime().init(request);
}
private AgentScopeReActRuntime fakeRuntime() { private AgentScopeReActRuntime fakeRuntime() {
return new AgentScopeReActRuntime(new FakeAgentScopeModelFactory(), new AgentScopeToolAdapter(), return new AgentScopeReActRuntime(new FakeAgentScopeModelFactory(), new AgentScopeToolAdapter(),
new AgentScopeKnowledgeAdapter(), new AgentScopeMemoryAdapter(), new AgentScopeSkillAdapter(), new AgentScopeKnowledgeAdapter(), new AgentScopeMemoryAdapter(), new AgentScopeSkillAdapter(),
@@ -1636,6 +1854,37 @@ public class AgentScopeStatefulRuntimeTest {
new AgentScopeMessageAdapter()); new AgentScopeMessageAdapter());
} }
/**
* 创建模型调用直接失败的运行时。
*
* @param error 模型异常
* @return 测试运行时
*/
private AgentScopeReActRuntime runtimeWithError(Throwable error) {
AgentScopeModelFactory modelFactory = new AgentScopeModelFactory() {
@Override
public Model create(AgentModelSpec modelSpec,
com.easyagents.agent.runtime.model.AgentGenerationOptions generationOptions) {
return new Model() {
@Override
public Flux<ChatResponse> stream(List<Msg> messages,
List<ToolSchema> toolSchemas,
GenerateOptions options) {
return Flux.error(error);
}
@Override
public String getModelName() {
return modelSpec == null ? "fake-model" : modelSpec.getModelName();
}
};
}
};
return new AgentScopeReActRuntime(modelFactory, new AgentScopeToolAdapter(),
new AgentScopeKnowledgeAdapter(), new AgentScopeMemoryAdapter(), new AgentScopeSkillAdapter(),
new AgentScopeMessageAdapter());
}
/** /**
* 创建每次模型调用仅返回下一条预设响应的运行时。 * 创建每次模型调用仅返回下一条预设响应的运行时。
* *

View File

@@ -0,0 +1,107 @@
package com.easyagents.agent.runtime.agentscope;
import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.AgentRuntimeExecutionContext;
import com.easyagents.agent.runtime.tool.AgentToolInvoker;
import com.easyagents.agent.runtime.tool.AgentToolResult;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import io.agentscope.core.message.ToolResultBlock;
import io.agentscope.core.message.ToolUseBlock;
import io.agentscope.core.tool.AgentTool;
import io.agentscope.core.tool.ToolCallParam;
import org.junit.Assert;
import org.junit.Test;
import java.lang.reflect.Proxy;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;
/**
* {@link AgentScopeToolAdapter} 回归测试。
*/
public class AgentScopeToolAdapterTest {
/**
* 验证 Reactor 工作线程执行工具时使用调用器所属类加载器,并在结束后恢复原上下文。
*/
@Test
public void shouldUseInvokerClassLoaderAndRestoreWorkerContext() {
ClassLoader parentClassLoader = AgentToolInvoker.class.getClassLoader();
ClassLoader invocationClassLoader = new ClassLoader(parentClassLoader) {
};
AtomicReference<ClassLoader> observedClassLoader = new AtomicReference<>();
AgentToolInvoker invoker = (AgentToolInvoker) Proxy.newProxyInstance(
invocationClassLoader,
new Class<?>[]{AgentToolInvoker.class},
(proxy, method, arguments) -> {
observedClassLoader.set(Thread.currentThread().getContextClassLoader());
return AgentToolResult.success("ok");
});
AgentTool tool = new AgentScopeToolAdapter().adapt(toolSpec(), invoker, executionContext());
Thread currentThread = Thread.currentThread();
ClassLoader originalClassLoader = currentThread.getContextClassLoader();
ClassLoader workerClassLoader = new ClassLoader(originalClassLoader) {
};
try {
currentThread.setContextClassLoader(workerClassLoader);
ToolResultBlock result = tool.callAsync(toolCall()).block();
Assert.assertNotNull(result);
Assert.assertSame(invocationClassLoader, observedClassLoader.get());
Assert.assertSame(workerClassLoader, currentThread.getContextClassLoader());
} finally {
currentThread.setContextClassLoader(originalClassLoader);
}
}
/**
* 创建测试工具声明。
*
* @return 工具声明
*/
private AgentToolSpec toolSpec() {
AgentToolSpec toolSpec = new AgentToolSpec();
toolSpec.setName("test_tool");
toolSpec.setDescription("Test tool");
toolSpec.setParametersSchema(Map.of(
"type", "object",
"properties", Map.of(),
"additionalProperties", false));
return toolSpec;
}
/**
* 创建测试运行上下文。
*
* @return 运行上下文
*/
private AgentRuntimeExecutionContext executionContext() {
AgentDefinition definition = new AgentDefinition();
definition.setAgentId("agent-1");
AgentRuntimeExecutionContext context = new AgentRuntimeExecutionContext();
context.setAgentDefinition(definition);
context.setRequestId("request-1");
context.setTraceId("trace-1");
context.setSessionId("session-1");
return context;
}
/**
* 创建测试工具调用参数。
*
* @return 工具调用参数
*/
private ToolCallParam toolCall() {
ToolUseBlock toolUseBlock = ToolUseBlock.builder()
.id("call-1")
.name("test_tool")
.input(Map.of())
.content("{}")
.build();
return ToolCallParam.builder()
.toolUseBlock(toolUseBlock)
.input(Map.of())
.build();
}
}

View File

@@ -72,6 +72,32 @@ public class HeuristicKnowledgeCitationMatcherTest {
Assert.assertTrue(references.isEmpty()); Assert.assertTrue(references.isEmpty());
} }
/**
* 长篇多主题汇总回答应按独立条目匹配引用,避免整篇答案稀释局部证据。
*/
@Test
public void shouldMatchReferencesFromLongMultiTopicAnswer() {
AgentKnowledgeReference checkIn = reference("faq-check-in",
"问题:最早入住酒店时间说明 答案最早入住时间为入住日当天下午14:00提前到店按房间状况安排。");
AgentKnowledgeReference luggage = reference("faq-luggage",
"问题:酒店寄存行李服务说明 答案离店客人通常可免费寄存2天第3天起收费。");
AgentKnowledgeReference unrelated = reference("faq-unrelated",
"问题:会员卡如何补办 答案:请携带本人证件到指定服务网点申请补办。");
String answer = "根据知识库整理,主要主题如下:\n"
+ "一、预订与支付官方渠道包括APP、微信小程序和客服电话。\n"
+ "二、入住与退房最早入住时间为下午14:00提前到店按房态安排。\n"
+ "三、酒店服务离店后行李通常可免费寄存2天第3天起收费。\n"
+ "四、会员商城:可使用彩虹如愿豆兑换商品。";
List<AgentKnowledgeReference> references = matcher.match(
answer, List.of(unrelated, checkIn, luggage));
Assert.assertEquals(2, references.size());
Assert.assertTrue(references.stream().anyMatch(reference -> "faq-check-in".equals(reference.getDocumentId())));
Assert.assertTrue(references.stream().anyMatch(reference -> "faq-luggage".equals(reference.getDocumentId())));
}
/** /**
* 空答案或空候选不应返回引用。 * 空答案或空候选不应返回引用。
*/ */

View File

@@ -0,0 +1,149 @@
package com.easyagents.agent.runtime.prompt;
import com.easyagents.agent.runtime.AgentDefinition;
import com.easyagents.agent.runtime.knowledge.AgentKnowledgeSpec;
import com.easyagents.agent.runtime.tool.AgentToolSpec;
import org.junit.Assert;
import org.junit.Test;
import java.util.List;
/**
* 测试运行时系统提示词组合器。
*/
public class SystemPromptComposerTest {
@Test
public void shouldKeepUserPromptWhenNoRuntimeProtocolIsRequired() {
AgentDefinition definition = definition(" user prompt ");
Assert.assertEquals(" user prompt ", SystemPromptComposer.compose(definition));
}
@Test
public void shouldReturnKnowledgeProtocolWhenUserPromptIsBlank() {
AgentDefinition definition = definition(" ");
definition.setKnowledgeSpecs(List.of(knowledge("faq")));
String prompt = SystemPromptComposer.compose(definition);
Assert.assertTrue(prompt.startsWith("Knowledge tool protocol:"));
Assert.assertFalse(prompt.startsWith("\n"));
}
@Test
public void shouldAppendKnowledgeProtocolAfterUserPrompt() {
AgentDefinition definition = definition("user prompt");
definition.setKnowledgeSpecs(List.of(knowledge("faq")));
String prompt = SystemPromptComposer.compose(definition);
Assert.assertTrue(prompt.startsWith("user prompt\n\nKnowledge tool protocol:"));
Assert.assertTrue(prompt.contains("prefer making one retrieval call"));
Assert.assertTrue(prompt.contains("reformulate the query and retry once"));
}
@Test
public void shouldAppendKnowledgeProtocolOnlyOnceForMultipleKnowledgeTools() {
AgentDefinition definition = definition("user prompt");
definition.setKnowledgeSpecs(List.of(knowledge("faq"), knowledge("policy")));
String prompt = SystemPromptComposer.compose(definition);
Assert.assertEquals(1, occurrences(prompt, "Knowledge tool protocol:"));
}
@Test
public void shouldSkipKnowledgeProtocolWhenToolCallingIsDisabled() {
AgentDefinition definition = definition("user prompt");
definition.setKnowledgeSpecs(List.of(knowledge("faq")));
definition.getExecutionOptions().setToolCallingEnabled(false);
Assert.assertEquals("user prompt", SystemPromptComposer.compose(definition));
}
@Test
public void shouldPreserveAsyncToolProtocolBehavior() {
AgentDefinition definition = definition(null);
definition.setToolSpecs(List.of(asyncTool()));
String prompt = SystemPromptComposer.compose(definition);
Assert.assertTrue(prompt.startsWith("Async tool protocol:"));
Assert.assertTrue(prompt.contains("These are internal execution phases."));
Assert.assertTrue(prompt.contains("call its submit sub-tool first with the user-provided arguments by default"));
Assert.assertTrue(prompt.contains("immediately call observe"));
}
@Test
public void shouldComposeUserKnowledgeAndAsyncProtocolsInStableOrder() {
AgentDefinition definition = definition("user prompt");
definition.setKnowledgeSpecs(List.of(knowledge("faq")));
definition.setToolSpecs(List.of(asyncTool()));
String prompt = SystemPromptComposer.compose(definition);
int userIndex = prompt.indexOf("user prompt");
int knowledgeIndex = prompt.indexOf("Knowledge tool protocol:");
int asyncIndex = prompt.indexOf("Async tool protocol:");
Assert.assertTrue(userIndex >= 0);
Assert.assertTrue(knowledgeIndex > userIndex);
Assert.assertTrue(asyncIndex > knowledgeIndex);
Assert.assertEquals(1, occurrences(prompt, "Knowledge tool protocol:"));
Assert.assertEquals(1, occurrences(prompt, "Async tool protocol:"));
}
/**
* 创建测试智能体定义。
*
* @param systemPrompt 用户系统提示词
* @return 智能体定义
*/
private AgentDefinition definition(String systemPrompt) {
AgentDefinition definition = new AgentDefinition();
definition.setSystemPrompt(systemPrompt);
return definition;
}
/**
* 创建测试知识库定义。
*
* @param runtimeName 知识库运行名
* @return 知识库定义
*/
private AgentKnowledgeSpec knowledge(String runtimeName) {
AgentKnowledgeSpec spec = new AgentKnowledgeSpec();
spec.setKnowledgeId("knowledge-" + runtimeName);
spec.setRuntimeName(runtimeName);
return spec;
}
/**
* 创建测试异步工具定义。
*
* @return 异步工具定义
*/
private AgentToolSpec asyncTool() {
AgentToolSpec toolSpec = new AgentToolSpec();
toolSpec.setName("demo_submit");
toolSpec.getMetadata().put("asyncTool", true);
return toolSpec;
}
/**
* 统计文本片段出现次数。
*
* @param value 待检查文本
* @param fragment 目标片段
* @return 出现次数
*/
private int occurrences(String value, String fragment) {
int count = 0;
int offset = 0;
while ((offset = value.indexOf(fragment, offset)) >= 0) {
count++;
offset += fragment.length();
}
return count;
}
}

View File

@@ -43,6 +43,11 @@ public class StoreOptions extends Metadata {
public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) { public void setEmbeddingOptions(EmbeddingOptions embeddingOptions) {
throw new IllegalStateException("Can not set embeddingOptions to the default instance."); 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; private EmbeddingOptions embeddingOptions = EmbeddingOptions.DEFAULT;
/**
* Optional upper bound for one store operation.
*/
private Long timeoutMillis;
public String getCollectionName() { public String getCollectionName() {
return collectionName; return collectionName;
@@ -111,6 +121,17 @@ public class StoreOptions extends Metadata {
this.embeddingOptions = embeddingOptions; 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) { public static StoreOptions ofCollectionName(String collectionName) {
StoreOptions storeOptions = new StoreOptions(); 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

@@ -2,6 +2,7 @@ package com.easyagents.document.core.async;
import com.easyagents.core.util.StringUtil; import com.easyagents.core.util.StringUtil;
import com.easyagents.document.core.exception.DocumentParseException; 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.ParseResponse;
import com.easyagents.document.core.entity.ParseTaskInfo; import com.easyagents.document.core.entity.ParseTaskInfo;
import com.easyagents.document.core.entity.ParseTaskStatus; import com.easyagents.document.core.entity.ParseTaskStatus;
@@ -135,7 +136,7 @@ public class DocumentAsyncTaskManager {
} }
DocumentAsyncTaskRecord record = repository.find(taskId); DocumentAsyncTaskRecord record = repository.find(taskId);
if (record == null) { if (record == null) {
throw new DocumentParseException("Document async task not found: " + taskId); throw new DocumentAsyncTaskNotFoundException(taskId);
} }
return record; 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.ParseResult;
import com.easyagents.document.core.entity.ParseTaskInfo; import com.easyagents.document.core.entity.ParseTaskInfo;
import com.easyagents.document.core.entity.ParseTaskStatus; import com.easyagents.document.core.entity.ParseTaskStatus;
import com.easyagents.document.core.exception.DocumentAsyncTaskNotFoundException;
import org.junit.Assert; import org.junit.Assert;
import org.junit.Test; import org.junit.Test;
@@ -18,6 +19,21 @@ import java.util.concurrent.Executor;
*/ */
public class DocumentAsyncTaskManagerTest { 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 @Test
public void shouldTrackTaskLifecycleAndResult() { public void shouldTrackTaskLifecycleAndResult() {
Executor directExecutor = new Executor() { 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.ParseResponse;
import com.easyagents.document.core.entity.ParseResult; import com.easyagents.document.core.entity.ParseResult;
import com.easyagents.document.core.entity.XlsxParseRequest; 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.core.support.AbstractAsyncDocumentParseService;
import com.easyagents.document.xlsx.XlsxDocumentProvider; import com.easyagents.document.xlsx.XlsxDocumentProvider;
import com.easyagents.document.xlsx.model.XlsxCellArtifact; 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 class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseService<XlsxParseRequest> implements XlsxDocumentProvider {
public static final String PROVIDER_NAME = "mineru"; 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 MineruProperties properties;
private final MineruClient client; private final MineruClient client;
@@ -145,6 +150,9 @@ public class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseSe
@Override @Override
protected ParseResponse doParse(XlsxParseRequest request, DocumentAsyncTaskUpdater updater) { protected ParseResponse doParse(XlsxParseRequest request, DocumentAsyncTaskUpdater updater) {
for (ParseFile file : request.getFiles()) {
validateXlsxContent(file);
}
ParseResponse response = new ParseResponse(); ParseResponse response = new ParseResponse();
List<ParseResult> results = new ArrayList<ParseResult>(); List<ParseResult> results = new ArrayList<ParseResult>();
String backend = null; String backend = null;
@@ -211,6 +219,50 @@ public class MineruXlsxDocumentParseService extends AbstractAsyncDocumentParseSe
return aggregate; 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, private SheetExtraction extractSheet(XSSFSheet sheet,
int sheetIndex, int sheetIndex,
DataFormatter formatter, 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.entity.XlsxParseRequest;
import com.easyagents.document.core.exception.DocumentParseException; import com.easyagents.document.core.exception.DocumentParseException;
import com.easyagents.document.xlsx.model.XlsxParseArtifact; 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.ss.usermodel.ClientAnchor;
import org.apache.poi.xssf.usermodel.XSSFDrawing; import org.apache.poi.xssf.usermodel.XSSFDrawing;
import org.apache.poi.xssf.usermodel.XSSFSheet; import org.apache.poi.xssf.usermodel.XSSFSheet;
@@ -119,6 +120,31 @@ public class MineruXlsxDocumentParseServiceTest {
Assert.assertEquals("image/jpeg", result.getImages().get(0).getMimeType()); 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 @Test
public void shouldAppendImageReferenceForImageOnlySheet() throws Exception { public void shouldAppendImageReferenceForImageOnlySheet() throws Exception {
RecordingClient client = new RecordingClient(defaultProperties()); RecordingClient client = new RecordingClient(defaultProperties());
@@ -255,6 +281,15 @@ public class MineruXlsxDocumentParseServiceTest {
return writeWorkbook(workbook); 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, private void addPicture(XSSFWorkbook workbook,
XSSFSheet sheet, XSSFSheet sheet,
int rowIndex, int rowIndex,

View File

@@ -0,0 +1,129 @@
package com.easyagents.federation.sql.adapter.jdbc;
import com.easyagents.federation.sql.api.FederationSqlEngine;
import com.easyagents.federation.sql.api.FederationSqlEngines;
import com.easyagents.federation.sql.api.SqlQueryCommand;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
import com.easyagents.federation.sql.source.FederationSourceDefinition;
import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceApplyOptions;
import com.easyagents.federation.sql.source.SourceId;
import com.mysql.cj.jdbc.MysqlDataSource;
import org.junit.Assert;
import org.junit.Assume;
import org.junit.Test;
import java.math.BigDecimal;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;
import java.util.Map;
/**
* 亿级数据集 100,000 行 Cursor 手工基准;默认跳过,必须显式启用。
*/
public class Jdbc100mCursorManualBenchmarkTest {
/**
* 从支付库顺序消费恰好 100,000 行,验证 Cursor 流式消费和资源上限。
*/
@Test
public void shouldConsumeOneHundredThousandRows() {
Assume.assumeTrue(Boolean.getBoolean("federation.100m.enabled"));
String host = System.getProperty("federation.100m.mysql.host", "127.0.0.1");
int port = Integer.getInteger("federation.100m.mysql.port", 33307);
String database = System.getProperty(
"federation.100m.mysql.database", "ef_bank_payment_perf");
String username = System.getProperty("federation.100m.mysql.username", "ef_bench_ro");
String password = System.getProperty("federation.100m.mysql.password");
Assert.assertNotNull("必须通过系统属性提供只读密码", password);
MysqlDataSource dataSource = new MysqlDataSource();
dataSource.setUrl("jdbc:mysql://" + host + ':' + port + '/' + database
+ "?useSSL=false&allowPublicKeyRetrieval=true&serverTimezone=Asia/Shanghai"
+ "&useCursorFetch=true&defaultFetchSize=1000");
dataSource.setUser(username);
dataSource.setPassword(password);
SourceId sourceId = new SourceId("payment-100m");
FederationSourceDefinition definition = new FederationSourceDefinition(
sourceId, 1, JdbcFederationSqlAdapterProvider.ADAPTER_ID,
List.of(new JdbcSchemaDefinition("MAIN", database, null)), Map.of());
FederationExecutionPolicy policy = new FederationExecutionPolicy(
6, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);
FederationQueryScopeDefinition scope = FederationQueryScopeDefinition.virtual(
"payment-cursor-100m", 1,
Map.of("PAY", FederationSourceBindingDefinition.of(sourceId, 1)),
"PAY", policy);
CursorResult nativeResult = consumeNative(dataSource);
CursorResult federationResult;
long federationStarted = System.nanoTime();
try (FederationSqlEngine engine = FederationSqlEngines.builder()
.dataSourceResolver(ignored -> FederationDataSourceHandles.shared(
dataSource, new RuntimeFingerprint("MySQL", "8.0", "JDBC", "8.4", "1")))
.adapter(new JdbcFederationSqlAdapterProvider())
.federationExecutionPolicy(policy)
.build()) {
engine.sources().apply(definition, SourceApplyOptions.prewarmNow());
String sql = "SELECT payment_id, amount FROM PAY.MAIN.payment_order "
+ "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
int rows = 0;
long amountCents = 0;
try (FederationResultCursor cursor = engine.query(
SqlQueryCommand.of(sql, scope, List.of()))) {
while (cursor.next()) {
rows++;
amountCents += ((BigDecimal) cursor.getObject(2))
.movePointRight(2).longValue();
}
}
federationResult = new CursorResult(rows, amountCents,
elapsedMillis(federationStarted));
}
Assert.assertEquals(100_000, nativeResult.rows());
Assert.assertEquals(100_000, federationResult.rows());
Assert.assertEquals(nativeResult.amountCents(), federationResult.amountCents());
Assert.assertTrue(federationResult.amountCents() > 0L);
System.out.println("BENCH_CURSOR_100K rows=" + federationResult.rows()
+ " nativeMillis=" + nativeResult.elapsedMillis()
+ " federationMillis=" + federationResult.elapsedMillis()
+ " amountCents=" + federationResult.amountCents());
}
private CursorResult consumeNative(MysqlDataSource dataSource) {
String sql = "SELECT payment_id, amount FROM payment_order "
+ "WHERE payment_id BETWEEN 1 AND 100000 ORDER BY payment_id";
long started = System.nanoTime();
int rows = 0;
long amountCents = 0;
try (Connection connection = dataSource.getConnection();
PreparedStatement statement = connection.prepareStatement(sql,
ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY)) {
statement.setFetchSize(1000);
try (ResultSet resultSet = statement.executeQuery()) {
while (resultSet.next()) {
rows++;
amountCents += resultSet.getBigDecimal(2)
.movePointRight(2).longValue();
}
}
} catch (SQLException exception) {
throw new IllegalStateException("原生 JDBC Cursor 基准执行失败", exception);
}
return new CursorResult(rows, amountCents, elapsedMillis(started));
}
private long elapsedMillis(long started) {
return (System.nanoTime() - started) / 1_000_000L;
}
private record CursorResult(int rows, long amountCents, long elapsedMillis) {
}
}

View File

@@ -27,6 +27,7 @@ import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceApplyOptions; import com.easyagents.federation.sql.source.SourceApplyOptions;
import com.easyagents.federation.sql.source.SourceId; import com.easyagents.federation.sql.source.SourceId;
import java.math.BigDecimal; import java.math.BigDecimal;
import java.sql.SQLException;
import java.time.Instant; import java.time.Instant;
import java.sql.Connection; import java.sql.Connection;
import java.sql.Statement; import java.sql.Statement;
@@ -34,6 +35,9 @@ import java.sql.Types;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference; import java.util.concurrent.atomic.AtomicReference;
import org.h2.jdbcx.JdbcDataSource; import org.h2.jdbcx.JdbcDataSource;
@@ -50,6 +54,13 @@ public class JdbcFederatedQueryEngineTest {
private static final SourceId SALES_SOURCE = new SourceId("sales-source"); private static final SourceId SALES_SOURCE = new SourceId("sales-source");
private static final SourceId BILLING_SOURCE = new SourceId("billing-source"); private static final SourceId BILLING_SOURCE = new SourceId("billing-source");
private static final SourceId REGION_SOURCE = new SourceId("region-source"); private static final SourceId REGION_SOURCE = new SourceId("region-source");
private static final AtomicInteger ACTIVE_FRAGMENTS = new AtomicInteger();
private static final AtomicInteger MAXIMUM_ACTIVE_FRAGMENTS = new AtomicInteger();
private static final AtomicInteger FAILING_TRACKED_VALUE = new AtomicInteger(
Integer.MIN_VALUE
);
private static final AtomicReference<CountDownLatch> FRAGMENT_START_BARRIER =
new AtomicReference<>(new CountDownLatch(0));
private JdbcDataSource sales; private JdbcDataSource sales;
private JdbcDataSource billing; private JdbcDataSource billing;
@@ -80,7 +91,10 @@ public class JdbcFederatedQueryEngineTest {
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 12:00:00.123456+08:00')", + "TIMESTAMP WITH TIME ZONE '2026-08-21 12:00:00.123456+08:00')",
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))", "CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
"INSERT INTO PRECISE_EVENT VALUES " "INSERT INTO PRECISE_EVENT VALUES "
+ "(1, TIMESTAMP '2026-08-21 12:00:00.123456')"); + "(1, TIMESTAMP '2026-08-21 12:00:00.123456')",
fragmentTrackingAlias(),
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(ID) AS TRACKED_KEY "
+ "FROM CUSTOMER WHERE ID <= 2");
execute(billing, execute(billing,
"CREATE TABLE ORDER_ITEM (ID INT PRIMARY KEY, CUSTOMER_ID INT NOT NULL, AMOUNT DECIMAL(12,2), CODE VARCHAR(64))", "CREATE TABLE ORDER_ITEM (ID INT PRIMARY KEY, CUSTOMER_ID INT NOT NULL, AMOUNT DECIMAL(12,2), CODE VARCHAR(64))",
"INSERT INTO ORDER_ITEM VALUES (10, 1, 30.00, 'Alice'), (11, 1, 20.00, 'Alice'), (12, 2, 80.00, 'Bob')", "INSERT INTO ORDER_ITEM VALUES (10, 1, 30.00, 'Alice'), (11, 1, 20.00, 'Alice'), (12, 2, 80.00, 'Bob')",
@@ -90,10 +104,16 @@ public class JdbcFederatedQueryEngineTest {
+ "TIMESTAMP WITH TIME ZONE '2026-08-21 06:00:00.123456+02:00')", + "TIMESTAMP WITH TIME ZONE '2026-08-21 06:00:00.123456+02:00')",
"CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))", "CREATE TABLE PRECISE_EVENT (ID INT PRIMARY KEY, EVENT_AT TIMESTAMP(6))",
"INSERT INTO PRECISE_EVENT VALUES " "INSERT INTO PRECISE_EVENT VALUES "
+ "(2, TIMESTAMP '2026-08-21 08:30:00.654321')"); + "(2, TIMESTAMP '2026-08-21 08:30:00.654321')",
fragmentTrackingAlias(),
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(ID) AS TRACKED_KEY "
+ "FROM ORDER_ITEM WHERE ID <= 11");
execute(region, execute(region,
"CREATE TABLE CUSTOMER_REGION (CUSTOMER_ID INT PRIMARY KEY, REGION VARCHAR(64))", "CREATE TABLE CUSTOMER_REGION (CUSTOMER_ID INT PRIMARY KEY, REGION VARCHAR(64))",
"INSERT INTO CUSTOMER_REGION VALUES (1, 'North'), (2, 'South'), (3, 'West')"); "INSERT INTO CUSTOMER_REGION VALUES (1, 'North'), (2, 'South'), (3, 'West')",
fragmentTrackingAlias(),
"CREATE VIEW TRACKED_VALUE AS SELECT TRACK_FRAGMENT(CUSTOMER_ID) AS TRACKED_KEY "
+ "FROM CUSTOMER_REGION WHERE CUSTOMER_ID <= 2");
RuntimeFingerprint fingerprint = new RuntimeFingerprint( RuntimeFingerprint fingerprint = new RuntimeFingerprint(
"H2", "2", "H2 JDBC Driver", "2", "1" "H2", "2", "H2 JDBC Driver", "2", "1"
@@ -356,6 +376,91 @@ public class JdbcFederatedQueryEngineTest {
); );
} }
/**
* 验证真实 JDBC Fragment 执行会重叠,且始终受查询级并发上限约束。
*/
@Test
public void shouldBoundConcurrentJdbcFragments() {
ACTIVE_FRAGMENTS.set(0);
MAXIMUM_ACTIVE_FRAGMENTS.set(0);
FRAGMENT_START_BARRIER.set(new CountDownLatch(2));
FederationQueryScopeDefinition threeSourceScope =
new FederationQueryScopeDefinition(
"tracked-fragments",
1,
Map.of(
"SALES", FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
"BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1),
"REGION", FederationSourceBindingDefinition.of(REGION_SOURCE, 1)
),
"SALES",
threeSourcePolicy()
);
try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
"SELECT COUNT(*) FROM SALES.APP.TRACKED_VALUE s "
+ "JOIN BILLING.APP.TRACKED_VALUE b ON s.TRACKED_KEY = b.TRACKED_KEY "
+ "JOIN REGION.APP.TRACKED_VALUE r ON s.TRACKED_KEY = r.TRACKED_KEY",
threeSourceScope,
List.of()
))) {
Assert.assertTrue(cursor.next());
Assert.assertEquals(8L, cursor.getObject(1));
Assert.assertFalse(cursor.next());
}
Assert.assertEquals(2, MAXIMUM_ACTIVE_FRAGMENTS.get());
Assert.assertEquals(0, ACTIVE_FRAGMENTS.get());
}
/**
* 验证任一并行 Fragment 失败后立即传播原始执行错误,不等待未启动任务超时。
*/
@Test
public void shouldFailFastWhenConcurrentFragmentFails() {
ACTIVE_FRAGMENTS.set(0);
MAXIMUM_ACTIVE_FRAGMENTS.set(0);
FRAGMENT_START_BARRIER.set(new CountDownLatch(0));
FAILING_TRACKED_VALUE.set(10);
FederationQueryScopeDefinition threeSourceScope =
new FederationQueryScopeDefinition(
"tracked-fragment-failure",
1,
Map.of(
"SALES", FederationSourceBindingDefinition.of(SALES_SOURCE, 1),
"BILLING", FederationSourceBindingDefinition.of(BILLING_SOURCE, 1),
"REGION", FederationSourceBindingDefinition.of(REGION_SOURCE, 1)
),
"SALES",
new FederationExecutionPolicy(
3,
8,
2,
100_000,
64L * 1024L * 1024L,
1_000
)
);
try (FederationResultCursor cursor = engine.query(SqlQueryCommand.of(
"SELECT COUNT(*) FROM SALES.APP.TRACKED_VALUE s "
+ "JOIN BILLING.APP.TRACKED_VALUE b ON s.TRACKED_KEY = b.TRACKED_KEY "
+ "JOIN REGION.APP.TRACKED_VALUE r ON s.TRACKED_KEY = r.TRACKED_KEY",
threeSourceScope,
List.of()
))) {
cursor.next();
Assert.fail("failing fragment should terminate the federated query");
} catch (FederationSqlException exception) {
Assert.assertEquals(
FederationSqlErrorCode.EXECUTION_FAILED,
exception.errorCode()
);
} finally {
FAILING_TRACKED_VALUE.set(Integer.MIN_VALUE);
}
}
/** /**
* 验证两个分片分别绑定原始查询中的动态参数。 * 验证两个分片分别绑定原始查询中的动态参数。
*/ */
@@ -717,6 +822,41 @@ public class JdbcFederatedQueryEngineTest {
return dataSource; return dataSource;
} }
/**
* H2 视图调用的并发跟踪函数。
*
* @param value 原始列值
* @return 固定键值
* @throws SQLException 分片未能在时限内重叠启动
*/
public static int trackFragment(int value) throws SQLException {
int active = ACTIVE_FRAGMENTS.incrementAndGet();
MAXIMUM_ACTIVE_FRAGMENTS.accumulateAndGet(active, Math::max);
CountDownLatch barrier = FRAGMENT_START_BARRIER.get();
barrier.countDown();
try {
if (value == FAILING_TRACKED_VALUE.get()) {
throw new SQLException("simulated federation fragment failure");
}
if (!barrier.await(2, TimeUnit.SECONDS)) {
throw new SQLException("federation fragments did not overlap");
}
Thread.sleep(100L);
return 1;
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
throw new SQLException("fragment tracking was interrupted", exception);
} finally {
ACTIVE_FRAGMENTS.decrementAndGet();
}
}
private static String fragmentTrackingAlias() {
return "CREATE ALIAS TRACK_FRAGMENT FOR \""
+ JdbcFederatedQueryEngineTest.class.getName()
+ ".trackFragment\"";
}
/** /**
* 根据物理源选择测试数据库。 * 根据物理源选择测试数据库。
* *

View File

@@ -0,0 +1,185 @@
package com.easyagents.federation.sql.adapter.jdbc;
import com.easyagents.federation.sql.api.FederationSqlEngine;
import com.easyagents.federation.sql.api.FederationSqlEngines;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.api.SqlQueryCommand;
import com.easyagents.federation.sql.compile.FederationSqlPlan;
import com.easyagents.federation.sql.compile.SqlCompileRequest;
import com.easyagents.federation.sql.execute.FederationResultCursor;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import com.easyagents.federation.sql.federation.FederationQueryMode;
import com.easyagents.federation.sql.federation.FederationQueryScopeDefinition;
import com.easyagents.federation.sql.federation.FederationSourceBindingDefinition;
import com.easyagents.federation.sql.source.FederationDataSourceHandles;
import com.easyagents.federation.sql.source.FederationSourceDefinition;
import com.easyagents.federation.sql.source.RuntimeFingerprint;
import com.easyagents.federation.sql.source.SourceApplyOptions;
import com.easyagents.federation.sql.source.SourceId;
import org.h2.jdbcx.JdbcDataSource;
import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;
import java.sql.Connection;
import java.sql.Statement;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
/**
* 四源和六源 JDBC 多阶段 Join 集成测试。
*/
public class JdbcSixSourceQueryEngineTest {
private static final FederationExecutionPolicy SIX_SOURCE_POLICY =
new FederationExecutionPolicy(6, 12, 2, 100_000,
64L * 1024L * 1024L, 60_000);
private final Map<SourceId, JdbcDataSource> dataSources = new LinkedHashMap<>();
private final Map<String, FederationSourceBindingDefinition> bindings = new LinkedHashMap<>();
private FederationSqlEngine engine;
/**
* 创建六个独立 H2 JDBC 数据源。
*/
@Before
public void setUp() throws Exception {
for (int index = 1; index <= 6; index++) {
SourceId sourceId = new SourceId("source-" + index);
JdbcDataSource dataSource = dataSource("six-source-" + index);
execute(dataSource,
"DROP TABLE IF EXISTS SUBJECT_METRIC",
"CREATE TABLE SUBJECT_METRIC (SUBJECT_ID INT PRIMARY KEY, METRIC_VALUE INT NOT NULL)",
"INSERT INTO SUBJECT_METRIC VALUES (1, " + (index * 10)
+ "), (2, " + (index * 100) + ")");
dataSources.put(sourceId, dataSource);
bindings.put("S" + index, FederationSourceBindingDefinition.of(sourceId, 1));
}
RuntimeFingerprint fingerprint = new RuntimeFingerprint(
"H2", "2", "H2 JDBC Driver", "2", "1");
engine = FederationSqlEngines.builder()
.dataSourceResolver(definition -> FederationDataSourceHandles.shared(
dataSources.get(definition.sourceId()), fingerprint))
.adapter(new JdbcFederationSqlAdapterProvider())
.federationExecutionPolicy(SIX_SOURCE_POLICY)
.maximumPlanCacheEntries(32)
.build();
for (SourceId sourceId : dataSources.keySet()) {
engine.sources().apply(definition(sourceId), SourceApplyOptions.prewarmNow());
}
}
/**
* 关闭 Engine。
*/
@After
public void tearDown() {
if (engine != null) {
engine.close();
}
}
/**
* 验证同一六源 Scope 中只引用四源时生成四个 Fragment。
*/
@Test
public void shouldCompileFourReferencedSources() {
FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
selectSql(4), scope(SIX_SOURCE_POLICY)));
Assert.assertEquals(FederationQueryMode.FEDERATED, plan.queryMode());
Assert.assertEquals(4, plan.referencedSources().size());
Assert.assertEquals(4, plan.fragments().size());
Assert.assertEquals(3, plan.joinOptimizations().size());
}
/**
* 验证六个独立 JDBC 数据源可完成五阶段 Join 并返回稳定结果。
*/
@Test
public void shouldExecuteJoinAcrossSixSources() {
String sql = selectSql(6);
FederationSqlPlan plan = engine.compile(SqlCompileRequest.of(
sql, scope(SIX_SOURCE_POLICY)));
Assert.assertEquals(6, plan.referencedSources().size());
Assert.assertEquals(6, plan.fragments().size());
Assert.assertEquals(5, plan.joinOptimizations().size());
try (FederationResultCursor cursor = engine.query(
SqlQueryCommand.of(sql, scope(SIX_SOURCE_POLICY), List.of()))) {
Assert.assertTrue(cursor.next());
Assert.assertEquals(1, cursor.getObject(1));
Assert.assertEquals(10, cursor.getObject(2));
Assert.assertEquals(60, cursor.getObject(7));
Assert.assertTrue(cursor.next());
Assert.assertEquals(2, cursor.getObject(1));
Assert.assertEquals(600, cursor.getObject(7));
Assert.assertFalse(cursor.next());
Assert.assertEquals(6, cursor.metrics().fragments().size());
}
}
/**
* 验证 Scope 收紧为四源时,六源 SQL 在编译期稳定拒绝。
*/
@Test
public void shouldRejectSixSourcesWhenScopeLimitIsFour() {
FederationExecutionPolicy fourSourcePolicy = new FederationExecutionPolicy(
4, 12, 2, 100_000, 64L * 1024L * 1024L, 60_000);
try {
engine.compile(SqlCompileRequest.of(selectSql(6), scope(fourSourcePolicy)));
Assert.fail("expected source limit rejection");
} catch (FederationSqlException exception) {
Assert.assertEquals(FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
exception.errorCode());
}
}
private FederationQueryScopeDefinition scope(FederationExecutionPolicy policy) {
return new FederationQueryScopeDefinition(
"six-source-scope", 1, bindings, "S1", policy);
}
private static String selectSql(int sourceCount) {
StringBuilder sql = new StringBuilder("SELECT s1.SUBJECT_ID");
for (int index = 1; index <= sourceCount; index++) {
sql.append(", s").append(index).append(".METRIC_VALUE");
}
sql.append(" FROM S1.APP.SUBJECT_METRIC s1");
for (int index = 2; index <= sourceCount; index++) {
sql.append(" JOIN S").append(index).append(".APP.SUBJECT_METRIC s")
.append(index).append(" ON s").append(index)
.append(".SUBJECT_ID = s1.SUBJECT_ID");
}
return sql.append(" ORDER BY s1.SUBJECT_ID").toString();
}
private static JdbcDataSource dataSource(String name) {
JdbcDataSource dataSource = new JdbcDataSource();
dataSource.setURL("jdbc:h2:mem:" + name + ";DB_CLOSE_DELAY=-1");
dataSource.setUser("sa");
dataSource.setPassword("");
return dataSource;
}
private static void execute(JdbcDataSource dataSource, String... statements)
throws Exception {
try (Connection connection = dataSource.getConnection();
Statement statement = connection.createStatement()) {
for (String sql : statements) {
statement.execute(sql);
}
}
}
private static FederationSourceDefinition definition(SourceId sourceId) {
return new FederationSourceDefinition(
sourceId, 1, JdbcFederationSqlAdapterProvider.ADAPTER_ID,
List.of(new JdbcSchemaDefinition("APP", null, "PUBLIC")), Map.of());
}
}

View File

@@ -70,6 +70,7 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
private final CalciteFederationSqlCompiler compiler; private final CalciteFederationSqlCompiler compiler;
private final AdapterFederationStatisticsProvider automaticStatisticsProvider; private final AdapterFederationStatisticsProvider automaticStatisticsProvider;
private final NodeMemoryAdmissionController nodeMemoryAdmission; private final NodeMemoryAdmissionController nodeMemoryAdmission;
private final FederationFragmentScheduler fragmentScheduler;
private final CalciteSqlCompleter completer = new CalciteSqlCompleter(); private final CalciteSqlCompleter completer = new CalciteSqlCompleter();
private final QueryCancellationRegistry cancellations = new QueryCancellationRegistry(); private final QueryCancellationRegistry cancellations = new QueryCancellationRegistry();
private final ScheduledThreadPoolExecutor deadlineScheduler = deadlineScheduler(); private final ScheduledThreadPoolExecutor deadlineScheduler = deadlineScheduler();
@@ -297,6 +298,7 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
this.nodeMemoryAdmission = new NodeMemoryAdmissionController( this.nodeMemoryAdmission = new NodeMemoryAdmissionController(
maximumNodeIntermediateBytes maximumNodeIntermediateBytes
); );
this.fragmentScheduler = new FederationFragmentScheduler(executionPolicy);
this.planCache = new BoundedPlanCache( this.planCache = new BoundedPlanCache(
maximumPlanCacheEntries, maximumPlanCacheEntries,
maximumConcurrentCompilations, maximumConcurrentCompilations,
@@ -650,7 +652,8 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
registration, registration,
compiler.effectivePolicy(plan.queryScope()), compiler.effectivePolicy(plan.queryScope()),
metrics, metrics,
queryDeadline queryDeadline,
fragmentScheduler
); );
queryCursor = new FederatedResultCursor( queryCursor = new FederatedResultCursor(
plan, plan,
@@ -1080,6 +1083,11 @@ public final class DefaultFederationSqlEngine implements FederationSqlEngine {
} catch (RuntimeException exception) { } catch (RuntimeException exception) {
failure = exception; failure = exception;
} }
try {
fragmentScheduler.close();
} catch (RuntimeException exception) {
failure = append(failure, exception);
}
try { try {
sourceManager.close(); sourceManager.close();
} catch (RuntimeException exception) { } catch (RuntimeException exception) {

View File

@@ -62,6 +62,7 @@ final class FederatedResultCursor
this.session = session; this.session = session;
this.metrics = metrics; this.metrics = metrics;
this.maximumRows = maximumRows; this.maximumRows = maximumRows;
session.prepareFragments();
this.enumerator = plan.localBindable() this.enumerator = plan.localBindable()
.bind(new FederationDataContext( .bind(new FederationDataContext(
plan.localRootSchema(), plan.localRootSchema(),

View File

@@ -20,14 +20,21 @@ import java.time.LocalDateTime;
import java.time.LocalTime; import java.time.LocalTime;
import java.time.OffsetDateTime; import java.time.OffsetDateTime;
import java.time.ZoneOffset; import java.time.ZoneOffset;
import java.util.ArrayDeque;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Set; import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import java.util.concurrent.Semaphore; import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.apache.calcite.avatica.util.ByteString; import org.apache.calcite.avatica.util.ByteString;
import org.apache.calcite.linq4j.AbstractEnumerable; import org.apache.calcite.linq4j.AbstractEnumerable;
import org.apache.calcite.linq4j.Enumerable; import org.apache.calcite.linq4j.Enumerable;
@@ -50,9 +57,19 @@ final class FederationExecutionSession implements AutoCloseable {
private final FederationExecutionPolicy policy; private final FederationExecutionPolicy policy;
private final FederationQueryMetricsTracker metrics; private final FederationQueryMetricsTracker metrics;
private final QueryDeadline deadline; private final QueryDeadline deadline;
private final FederationFragmentScheduler fragmentScheduler;
private final Semaphore fragmentSlots; private final Semaphore fragmentSlots;
private final Set<FederationResultCursor> openFragmentCursors = ConcurrentHashMap.newKeySet(); private final Set<FederationResultCursor> openFragmentCursors = ConcurrentHashMap.newKeySet();
private final Set<Future<?>> fragmentTasks = ConcurrentHashMap.newKeySet();
private final Map<String, CompletableFuture<PreparedFragment>> fragmentPreparations =
new ConcurrentHashMap<>();
private final Object preparationLock = new Object();
private final ArrayDeque<CompiledFederationFragment> pendingPreparations =
new ArrayDeque<>();
private final AtomicBoolean closed = new AtomicBoolean(); private final AtomicBoolean closed = new AtomicBoolean();
private final AtomicReference<Throwable> preparationFailure = new AtomicReference<>();
private int runningPreparations;
private boolean preparationFailed;
/** /**
* 创建查询执行会话。 * 创建查询执行会话。
@@ -64,6 +81,7 @@ final class FederationExecutionSession implements AutoCloseable {
* @param policy 有效资源策略 * @param policy 有效资源策略
* @param metrics 指标跟踪器 * @param metrics 指标跟踪器
* @param deadline 请求级统一截止时间 * @param deadline 请求级统一截止时间
* @param fragmentScheduler Engine 级有界分片调度器
*/ */
FederationExecutionSession( FederationExecutionSession(
DefaultFederationSqlPlan plan, DefaultFederationSqlPlan plan,
@@ -72,7 +90,8 @@ final class FederationExecutionSession implements AutoCloseable {
QueryCancellationRegistry.QueryRegistration registration, QueryCancellationRegistry.QueryRegistration registration,
FederationExecutionPolicy policy, FederationExecutionPolicy policy,
FederationQueryMetricsTracker metrics, FederationQueryMetricsTracker metrics,
QueryDeadline deadline QueryDeadline deadline,
FederationFragmentScheduler fragmentScheduler
) { ) {
this.plan = plan; this.plan = plan;
this.context = context; this.context = context;
@@ -81,9 +100,41 @@ final class FederationExecutionSession implements AutoCloseable {
this.policy = policy; this.policy = policy;
this.metrics = metrics; this.metrics = metrics;
this.deadline = deadline; this.deadline = deadline;
this.fragmentScheduler = fragmentScheduler;
this.fragmentSlots = new Semaphore(policy.maximumConcurrentFragments(), true); this.fragmentSlots = new Semaphore(policy.maximumConcurrentFragments(), true);
} }
/**
* 异步启动多源物理分片,在查询级中间结果预算内完成有界物化。
*/
void prepareFragments() {
if (plan.fragments().size() < 2 || policy.maximumConcurrentFragments() < 2) {
return;
}
try {
for (FederationFragmentPlan fragment : plan.fragments()) {
CompiledFederationFragment compiled = plan.compiledFragments().get(
fragment.fragmentId()
);
if (compiled == null) {
throw new FederationSqlException(
FederationSqlErrorCode.EXECUTION_FAILED,
"compiled fragment is missing: " + fragment.fragmentId()
);
}
CompletableFuture<PreparedFragment> preparation = new CompletableFuture<>();
fragmentPreparations.put(fragment.fragmentId(), preparation);
synchronized (preparationLock) {
pendingPreparations.addLast(compiled);
}
}
schedulePreparations();
} catch (RuntimeException | Error exception) {
cancelPreparation();
throw exception;
}
}
/** /**
* 为 Calcite ScannableTable 创建一个按需执行物理分片的 Enumerable。 * 为 Calcite ScannableTable 创建一个按需执行物理分片的 Enumerable。
* *
@@ -101,6 +152,13 @@ final class FederationExecutionSession implements AutoCloseable {
return new AbstractEnumerable<>() { return new AbstractEnumerable<>() {
@Override @Override
public Enumerator<Object[]> enumerator() { public Enumerator<Object[]> enumerator() {
CompletableFuture<PreparedFragment> preparation =
fragmentPreparations.get(fragmentId);
if (preparation != null) {
PreparedFragment prepared = awaitPrepared(preparation);
fragmentPreparations.remove(fragmentId, preparation);
return prepared.enumerator();
}
return openFragment(compiled); return openFragment(compiled);
} }
}; };
@@ -138,6 +196,7 @@ final class FederationExecutionSession implements AutoCloseable {
} }
registration.ensureNotCancelled(); registration.ensureNotCancelled();
} }
throwPreparationFailure();
deadline.ensureAllowed(); deadline.ensureAllowed();
} }
@@ -166,9 +225,189 @@ final class FederationExecutionSession implements AutoCloseable {
private Enumerator<Object[]> openFragment(CompiledFederationFragment compiled) { private Enumerator<Object[]> openFragment(CompiledFederationFragment compiled) {
ensureExecutionAllowed(); ensureExecutionAllowed();
acquireFragmentSlot(); acquireFragmentSlot();
try {
return new FragmentEnumerator(compiled, createFragmentCursor(compiled), true);
} catch (RuntimeException exception) {
fragmentSlots.release();
throw exception;
}
}
private void schedulePreparations() {
List<PreparationTask> tasks = new ArrayList<>();
synchronized (preparationLock) {
while (!closed.get()
&& !preparationFailed
&& runningPreparations < policy.maximumConcurrentFragments()
&& !pendingPreparations.isEmpty()) {
CompiledFederationFragment compiled = pendingPreparations.removeFirst();
CompletableFuture<PreparedFragment> preparation =
fragmentPreparations.get(
compiled.plan().fragmentId()
);
if (preparation == null || preparation.isCancelled()) {
continue;
}
runningPreparations++;
tasks.add(new PreparationTask(compiled, preparation));
}
}
for (PreparationTask task : tasks) {
submitPreparation(task);
}
}
private void submitPreparation(PreparationTask preparationTask) {
FutureTask<Void> task = new FutureTask<>(() -> {
runPreparation(preparationTask);
return null;
}) {
@Override
protected void done() {
fragmentTasks.remove(this);
}
};
fragmentTasks.add(task);
try {
fragmentScheduler.execute(task);
} catch (RuntimeException | Error exception) {
fragmentTasks.remove(task);
task.cancel(true);
preparationTask.preparation().completeExceptionally(exception);
failPreparations(exception);
preparationFinished();
}
}
private void runPreparation(PreparationTask task) {
try {
task.preparation().complete(prepareFragment(task.compiled()));
} catch (RuntimeException | Error exception) {
task.preparation().completeExceptionally(exception);
failPreparations(exception);
} finally {
preparationFinished();
}
}
private void preparationFinished() {
boolean scheduleNext;
synchronized (preparationLock) {
runningPreparations--;
scheduleNext = !closed.get() && !preparationFailed;
}
if (scheduleNext) {
schedulePreparations();
}
}
private void failPreparations(Throwable failure) {
if (!preparationFailure.compareAndSet(null, failure)) {
return;
}
synchronized (preparationLock) {
preparationFailed = true;
pendingPreparations.clear();
}
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
preparation.completeExceptionally(failure);
}
registration.terminateStatementsForCleanup();
}
private void throwPreparationFailure() {
Throwable failure = preparationFailure.get();
if (failure == null) {
return;
}
if (failure instanceof RuntimeException runtimeException) {
throw runtimeException;
}
if (failure instanceof Error error) {
throw error;
}
throw new FederationSqlException(
FederationSqlErrorCode.EXECUTION_FAILED,
"federation fragment preparation failed",
failure
);
}
private PreparedFragment awaitPrepared(Future<PreparedFragment> preparation) {
try {
while (true) {
ensureExecutionAllowed();
long remaining = deadline.remainingNanos();
if (remaining <= 0L) {
ensureExecutionAllowed();
}
try {
return preparation.get(
Math.min(remaining, TimeUnit.MILLISECONDS.toNanos(50)),
TimeUnit.NANOSECONDS
);
} catch (TimeoutException ignored) {
// 周期性复核共享 Deadline 和取消状态。
}
}
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
cancelPreparation();
throw new FederationSqlException(
FederationSqlErrorCode.EXECUTION_FAILED,
"waiting for federation fragment preparation was interrupted",
exception
);
} catch (ExecutionException exception) {
cancelPreparation();
Throwable cause = exception.getCause();
if (cause instanceof RuntimeException runtimeException) {
throw runtimeException;
}
if (cause instanceof Error error) {
throw error;
}
throw new FederationSqlException(
FederationSqlErrorCode.EXECUTION_FAILED,
"federation fragment preparation failed",
cause
);
} catch (RuntimeException | Error exception) {
cancelPreparation();
throw exception;
}
}
private PreparedFragment prepareFragment(CompiledFederationFragment compiled) {
ensureExecutionAllowed();
FederationResultCursor cursor = null;
try {
cursor = createFragmentCursor(compiled);
List<Object[]> rows = new ArrayList<>();
while (true) {
ensureExecutionAllowed();
if (!cursor.next()) {
releaseFragment(compiled, cursor, true, false);
cursor = null;
return new PreparedFragment(rows);
}
Object[] row = normalizeRow(cursor, compiled.rowType());
recordFragmentIntermediate(compiled, row);
rows.add(row);
}
} catch (RuntimeException | Error exception) {
if (cursor != null) {
releaseFragment(compiled, cursor, false, false);
}
throw exception;
}
}
private FederationResultCursor createFragmentCursor(
CompiledFederationFragment compiled
) {
FederationFragmentPlan fragment = compiled.plan(); FederationFragmentPlan fragment = compiled.plan();
SourceRuntime runtime = runtimeSnapshot.runtime(fragment.bindingName()); SourceRuntime runtime = runtimeSnapshot.runtime(fragment.bindingName());
try {
FederationResultCursor cursor = runtime.adapter().fragmentExecutor().execute( FederationResultCursor cursor = runtime.adapter().fragmentExecutor().execute(
new FederationFragmentExecutionContext( new FederationFragmentExecutionContext(
context.queryId(), context.queryId(),
@@ -191,13 +430,55 @@ final class FederationExecutionSession implements AutoCloseable {
); );
} }
openFragmentCursors.add(cursor); openFragmentCursors.add(cursor);
return new FragmentEnumerator(compiled, cursor); try {
} catch (RuntimeException exception) { ensureExecutionAllowed();
fragmentSlots.release(); return cursor;
} catch (RuntimeException | Error exception) {
openFragmentCursors.remove(cursor);
closeQuietly(cursor);
throw exception; throw exception;
} }
} }
private void recordFragmentIntermediate(
CompiledFederationFragment compiled,
Object[] row
) {
FederationQueryMetricsTracker.IntermediateUsage usage =
metrics.recordIntermediate(compiled.plan().fragmentId(), row);
if (usage.rows() > policy.maximumIntermediateRows()
|| usage.bytes() > policy.maximumIntermediateBytes()) {
throw new FederationSqlException(
FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
"federation intermediate result limit was exceeded"
);
}
}
private void cancelPreparation() {
synchronized (preparationLock) {
preparationFailed = true;
pendingPreparations.clear();
}
for (Future<?> task : List.copyOf(fragmentTasks)) {
task.cancel(true);
}
fragmentTasks.clear();
for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
openFragmentCursors.remove(cursor);
try {
cursor.close();
} catch (RuntimeException ignored) {
// 保留触发取消的原始分片异常Statement 取消通道继续收口剩余资源。
}
}
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
preparation.cancel(true);
}
fragmentPreparations.clear();
registration.terminateStatementsForCleanup();
}
private void acquireFragmentSlot() { private void acquireFragmentSlot() {
while (true) { while (true) {
ensureExecutionAllowed(); ensureExecutionAllowed();
@@ -259,6 +540,21 @@ final class FederationExecutionSession implements AutoCloseable {
if (!closed.compareAndSet(false, true)) { if (!closed.compareAndSet(false, true)) {
return; return;
} }
boolean terminationRequired = !fragmentTasks.isEmpty()
|| !openFragmentCursors.isEmpty();
synchronized (preparationLock) {
terminationRequired = terminationRequired || runningPreparations > 0;
preparationFailed = true;
pendingPreparations.clear();
}
for (Future<?> task : List.copyOf(fragmentTasks)) {
task.cancel(true);
}
fragmentTasks.clear();
for (CompletableFuture<PreparedFragment> preparation : fragmentPreparations.values()) {
preparation.cancel(true);
}
fragmentPreparations.clear();
RuntimeException failure = null; RuntimeException failure = null;
for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) { for (FederationResultCursor cursor : List.copyOf(openFragmentCursors)) {
try { try {
@@ -273,6 +569,9 @@ final class FederationExecutionSession implements AutoCloseable {
openFragmentCursors.remove(cursor); openFragmentCursors.remove(cursor);
} }
} }
if (terminationRequired) {
registration.terminateStatementsForCleanup();
}
if (failure != null) { if (failure != null) {
throw failure; throw failure;
} }
@@ -282,15 +581,18 @@ final class FederationExecutionSession implements AutoCloseable {
private final CompiledFederationFragment compiled; private final CompiledFederationFragment compiled;
private final FederationResultCursor cursor; private final FederationResultCursor cursor;
private final boolean fragmentSlotHeld;
private final AtomicBoolean released = new AtomicBoolean(); private final AtomicBoolean released = new AtomicBoolean();
private Object[] current; private Object[] current;
private FragmentEnumerator( private FragmentEnumerator(
CompiledFederationFragment compiled, CompiledFederationFragment compiled,
FederationResultCursor cursor FederationResultCursor cursor,
boolean fragmentSlotHeld
) { ) {
this.compiled = compiled; this.compiled = compiled;
this.cursor = cursor; this.cursor = cursor;
this.fragmentSlotHeld = fragmentSlotHeld;
} }
/** {@inheritDoc} */ /** {@inheritDoc} */
@@ -309,15 +611,7 @@ final class FederationExecutionSession implements AutoCloseable {
return false; return false;
} }
Object[] row = normalizeRow(cursor, compiled.rowType()); Object[] row = normalizeRow(cursor, compiled.rowType());
FederationQueryMetricsTracker.IntermediateUsage usage = recordFragmentIntermediate(compiled, row);
metrics.recordIntermediate(compiled.plan().fragmentId(), row);
if (usage.rows() > policy.maximumIntermediateRows()
|| usage.bytes() > policy.maximumIntermediateBytes()) {
throw new FederationSqlException(
FederationSqlErrorCode.FEDERATION_RESOURCE_LIMIT_EXCEEDED,
"federation intermediate result limit was exceeded"
);
}
current = row; current = row;
return true; return true;
} catch (RuntimeException exception) { } catch (RuntimeException exception) {
@@ -346,7 +640,9 @@ final class FederationExecutionSession implements AutoCloseable {
try { try {
cursor.close(); cursor.close();
} finally { } finally {
if (fragmentSlotHeld) {
fragmentSlots.release(); fragmentSlots.release();
}
if (exhausted) { if (exhausted) {
metrics.finishFragment(compiled.plan().fragmentId()); metrics.finishFragment(compiled.plan().fragmentId());
} }
@@ -354,6 +650,78 @@ final class FederationExecutionSession implements AutoCloseable {
} }
} }
private static final class PreparedFragment {
private final List<Object[]> rows;
private PreparedFragment(List<Object[]> rows) {
this.rows = rows;
}
private Enumerator<Object[]> enumerator() {
return new Enumerator<>() {
private int index = -1;
private Object[] current;
@Override
public Object[] current() {
return current;
}
@Override
public boolean moveNext() {
int next = index + 1;
if (next >= rows.size()) {
current = null;
return false;
}
index = next;
current = rows.set(index, null);
return true;
}
@Override
public void reset() {
throw new UnsupportedOperationException(
"fragment cursor cannot be reset"
);
}
@Override
public void close() {
rows.clear();
current = null;
}
};
}
}
private record PreparationTask(
CompiledFederationFragment compiled,
CompletableFuture<PreparedFragment> preparation
) {
}
private void releaseFragment(
CompiledFederationFragment compiled,
FederationResultCursor cursor,
boolean exhausted,
boolean fragmentSlotHeld
) {
openFragmentCursors.remove(cursor);
try {
cursor.close();
} finally {
if (fragmentSlotHeld) {
fragmentSlots.release();
}
if (exhausted) {
metrics.finishFragment(compiled.plan().fragmentId());
}
}
}
private static Object[] normalizeRow( private static Object[] normalizeRow(
FederationResultCursor cursor, FederationResultCursor cursor,
RelDataType rowType RelDataType rowType

View File

@@ -0,0 +1,118 @@
package com.easyagents.federation.sql.runtime;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import java.util.concurrent.ArrayBlockingQueue;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
/**
* Engine 级有界分片工作线程池;查询级并发度仍由有效执行策略单独控制。
*/
final class FederationFragmentScheduler implements Executor, AutoCloseable {
private static final int MAXIMUM_WORKERS = 32;
private static final int MINIMUM_QUEUE_CAPACITY = 64;
private static final int QUEUE_CAPACITY_PER_WORKER = 8;
private final ThreadPoolExecutor executor;
private final AtomicBoolean closed = new AtomicBoolean();
/**
* 创建 Engine 独占的有界分片调度器。
*
* @param policy Engine 联邦执行硬上限
*/
FederationFragmentScheduler(FederationExecutionPolicy policy) {
int availableProcessors = Math.max(1, Runtime.getRuntime().availableProcessors());
int workers = Math.min(
MAXIMUM_WORKERS,
Math.max(policy.maximumConcurrentFragments(), availableProcessors)
);
int queueCapacity = Math.max(
MINIMUM_QUEUE_CAPACITY,
workers * QUEUE_CAPACITY_PER_WORKER
);
this.executor = executor(workers, queueCapacity);
}
FederationFragmentScheduler(int workers, int queueCapacity) {
if (workers <= 0) {
throw new IllegalArgumentException("workers must be positive");
}
if (queueCapacity <= 0) {
throw new IllegalArgumentException("queueCapacity must be positive");
}
this.executor = executor(workers, queueCapacity);
}
private static ThreadPoolExecutor executor(int workers, int queueCapacity) {
AtomicInteger threadSequence = new AtomicInteger();
return new ThreadPoolExecutor(
workers,
workers,
0L,
TimeUnit.MILLISECONDS,
new ArrayBlockingQueue<>(queueCapacity),
runnable -> {
Thread thread = new Thread(
runnable,
"easy-agents-federation-fragment-" + threadSequence.incrementAndGet()
);
thread.setDaemon(true);
return thread;
},
new ThreadPoolExecutor.AbortPolicy()
);
}
/**
* 提交分片任务;队列饱和时明确拒绝,禁止在调用线程绕过 Worker 上限执行 JDBC。
*
* @param command 分片任务
*/
@Override
public void execute(Runnable command) {
if (closed.get()) {
throw engineClosed();
}
try {
executor.execute(command);
} catch (RejectedExecutionException exception) {
if (closed.get() || executor.isShutdown()) {
throw engineClosed();
}
throw schedulerOverloaded(exception);
}
}
/**
* 中断排队任务并关闭工作线程。
*/
@Override
public void close() {
if (closed.compareAndSet(false, true)) {
executor.shutdownNow();
}
}
private static FederationSqlException engineClosed() {
return new FederationSqlException(
FederationSqlErrorCode.ENGINE_CLOSED,
"federation fragment scheduler is closed"
);
}
private static FederationSqlException schedulerOverloaded(Throwable cause) {
return new FederationSqlException(
FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
"federation fragment scheduler is saturated",
cause
);
}
}

View File

@@ -301,6 +301,13 @@ final class QueryCancellationRegistry implements AutoCloseable {
return state.cancellationRequested(); return state.cancellationRequested();
} }
/**
* 不改变查询终态,仅终止当前已登记的 Statement用于提前关闭异步分片。
*/
void terminateStatementsForCleanup() {
state.terminateStatements();
}
/** /**
* 已取消时立即抛出稳定错误。 * 已取消时立即抛出稳定错误。
*/ */
@@ -485,20 +492,26 @@ final class QueryCancellationRegistry implements AutoCloseable {
} }
private void terminate(QueryId queryId) { private void terminate(QueryId queryId) {
if (isClosed()) {
return;
}
SQLException failure = null; SQLException failure = null;
if (cancelIssued.compareAndSet(false, true)) { if (cancelIssued.compareAndSet(false, true)) {
try { try {
statement.cancel(); statement.cancel();
} catch (SQLException exception) { } catch (SQLException exception) {
if (!isClosed()) {
cancelIssued.set(false); cancelIssued.set(false);
failure = exception; failure = exception;
} }
} }
}
// close() 同时覆盖 register 与 executeQuery 之间的 JDBC 取消空窗。 // close() 同时覆盖 register 与 executeQuery 之间的 JDBC 取消空窗。
if (closeIssued.compareAndSet(false, true)) { if (closeIssued.compareAndSet(false, true)) {
try { try {
statement.close(); statement.close();
} catch (SQLException exception) { } catch (SQLException exception) {
if (!isClosed()) {
closeIssued.set(false); closeIssued.set(false);
if (failure == null) { if (failure == null) {
failure = exception; failure = exception;
@@ -507,6 +520,7 @@ final class QueryCancellationRegistry implements AutoCloseable {
} }
} }
} }
}
if (failure != null) { if (failure != null) {
throw new FederationSqlException( throw new FederationSqlException(
FederationSqlErrorCode.EXECUTION_FAILED, FederationSqlErrorCode.EXECUTION_FAILED,
@@ -515,6 +529,14 @@ final class QueryCancellationRegistry implements AutoCloseable {
); );
} }
} }
private boolean isClosed() {
try {
return statement.isClosed();
} catch (SQLException ignored) {
return false;
}
}
} }
private static FederationSqlException cancelled(QueryId queryId, Throwable cause) { private static FederationSqlException cancelled(QueryId queryId, Throwable cause) {

View File

@@ -0,0 +1,118 @@
package com.easyagents.federation.sql.runtime;
import com.easyagents.federation.sql.api.FederationSqlErrorCode;
import com.easyagents.federation.sql.api.FederationSqlException;
import com.easyagents.federation.sql.federation.FederationExecutionPolicy;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.junit.Assert;
import org.junit.Test;
/**
* Engine 级有界分片调度器测试。
*/
public class FederationFragmentSchedulerTest {
/**
* 验证两个分片任务能够真实重叠执行,且调度器关闭后拒绝新任务。
*
* @throws Exception 等待任务失败
*/
@Test
public void shouldExecuteFragmentTasksConcurrentlyAndRejectAfterClose() throws Exception {
FederationFragmentScheduler scheduler = new FederationFragmentScheduler(
FederationExecutionPolicy.basic()
);
CountDownLatch started = new CountDownLatch(2);
CountDownLatch release = new CountDownLatch(1);
AtomicInteger active = new AtomicInteger();
AtomicInteger maximumActive = new AtomicInteger();
ExecutorCompletionService<Void> completions = new ExecutorCompletionService<>(scheduler);
try {
Future<Void> first = completions.submit(
() -> runBlockingTask(started, release, active, maximumActive)
);
Future<Void> second = completions.submit(
() -> runBlockingTask(started, release, active, maximumActive)
);
Assert.assertTrue(started.await(2, TimeUnit.SECONDS));
Assert.assertEquals(2, maximumActive.get());
release.countDown();
first.get(2, TimeUnit.SECONDS);
second.get(2, TimeUnit.SECONDS);
} finally {
release.countDown();
scheduler.close();
}
try {
scheduler.execute(() -> { });
Assert.fail("closed scheduler should reject new tasks");
} catch (FederationSqlException exception) {
Assert.assertEquals(FederationSqlErrorCode.ENGINE_CLOSED, exception.errorCode());
}
}
/**
* 验证队列饱和时明确拒绝,且不会在提交线程绕过 Worker 上限执行任务。
*
* @throws Exception 等待任务失败
*/
@Test
public void shouldRejectSaturationWithoutRunningTaskOnCaller() throws Exception {
FederationFragmentScheduler scheduler = new FederationFragmentScheduler(1, 1);
CountDownLatch firstStarted = new CountDownLatch(1);
CountDownLatch releaseFirst = new CountDownLatch(1);
CountDownLatch queuedCompleted = new CountDownLatch(1);
AtomicBoolean rejectedTaskRan = new AtomicBoolean();
try {
scheduler.execute(() -> {
firstStarted.countDown();
try {
releaseFirst.await(2, TimeUnit.SECONDS);
} catch (InterruptedException exception) {
Thread.currentThread().interrupt();
}
});
Assert.assertTrue(firstStarted.await(2, TimeUnit.SECONDS));
scheduler.execute(queuedCompleted::countDown);
try {
scheduler.execute(() -> rejectedTaskRan.set(true));
Assert.fail("saturated scheduler should reject the task");
} catch (FederationSqlException exception) {
Assert.assertEquals(
FederationSqlErrorCode.QUERY_ADMISSION_TIMEOUT,
exception.errorCode()
);
}
Assert.assertFalse(rejectedTaskRan.get());
} finally {
releaseFirst.countDown();
Assert.assertTrue(queuedCompleted.await(2, TimeUnit.SECONDS));
scheduler.close();
}
}
private static Void runBlockingTask(
CountDownLatch started,
CountDownLatch release,
AtomicInteger active,
AtomicInteger maximumActive
) throws InterruptedException {
int current = active.incrementAndGet();
maximumActive.accumulateAndGet(current, Math::max);
started.countDown();
try {
Assert.assertTrue(release.await(2, TimeUnit.SECONDS));
return null;
} finally {
active.decrementAndGet();
}
}
}

View File

@@ -634,7 +634,8 @@ public class Chain {
NodeStateField NodeStateField
.EXECUTION_ATTEMPT_KEY); .EXECUTION_ATTEMPT_KEY);
} }
if (node.getCondition() == null) { if (node.getJoinMode() == NodeJoinMode.ANY
&& node.getCondition() == null) {
s.recordTrigger(triggerEdgeId); s.recordTrigger(triggerEdgeId);
fields.add(NodeStateField.TRIGGER_COUNT); fields.add(NodeStateField.TRIGGER_COUNT);
fields.add(NodeStateField.TRIGGER_EDGE_IDS); fields.add(NodeStateField.TRIGGER_EDGE_IDS);
@@ -795,7 +796,7 @@ public class Chain {
private boolean shouldSkipNode(Node node, String edgeId) { private boolean shouldSkipNode(Node node, String edgeId) {
NodeCondition condition = node.getCondition(); NodeCondition condition = node.getCondition();
if (condition == null) { if (node.getJoinMode() == NodeJoinMode.ANY && condition == null) {
return false; return false;
} }
return executeWithLock(stateInstanceId, 10, TimeUnit.SECONDS, () -> { return executeWithLock(stateInstanceId, 10, TimeUnit.SECONDS, () -> {
@@ -805,7 +806,11 @@ public class Chain {
}); });
Map<String, Object> prevResult = Collections.emptyMap(); Map<String, Object> prevResult = Collections.emptyMap();
boolean shouldSkipNode = !condition.check(this, newState, prevResult); boolean joinPending = node.getJoinMode() == NodeJoinMode.ALL
&& !newState.isUpstreamFullyExecuted();
boolean shouldSkipNode = joinPending
|| (condition != null
&& !condition.check(this, newState, prevResult));
if (shouldSkipNode) { if (shouldSkipNode) {
updateStateSafely(state -> { updateStateSafely(state -> {
return state.addUncheckedNodeId(node.id) return state.addUncheckedNodeId(node.id)
@@ -1746,16 +1751,91 @@ public class Chain {
stateInstanceId, stateInstanceId,
10L, 10L,
TimeUnit.SECONDS, TimeUnit.SECONDS,
() -> { () -> resumeSuspended(variables));
ChainState current =
chainStateRepository.load(stateInstanceId);
if (current == null
|| current.getStatus() != ChainStatus.SUSPEND) {
return false;
} }
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; return true;
}); }
if (value instanceof String text) {
return text.trim().isEmpty();
}
return value instanceof Collection<?> collection
&& collection.isEmpty();
} }
/** /**
@@ -1772,36 +1852,51 @@ public class Chain {
* *
* @param variables 恢复时注入的变量 * @param variables 恢复时注入的变量
*/ */
private void resumeSuspended(Map<String, Object> variables) { private boolean resumeSuspended(Map<String, Object> variables) {
ChainState newState = updateStateSafely(state -> { AtomicBoolean resumed = new AtomicBoolean(false);
if (variables != null) { AtomicReference<Set<String>> suspendedNodeIds =
state.getMemory().putAll(variables); new AtomicReference<>(Collections.emptySet());
return EnumSet.of(ChainStateField.MEMORY); updateStateSafely(state -> {
} else { resumed.set(false);
suspendedNodeIds.set(Collections.emptySet());
if (state.getStatus() != ChainStatus.SUSPEND) {
return null; return null;
} }
}); validateResumeVariables(state, variables);
if (state.getSuspendNodeIds() != null) {
notifyEvent(new ChainResumeEvent(this, variables)); suspendedNodeIds.set(
setStatusAndNotifyEvent(ChainStatus.RUNNING); new LinkedHashSet<>(state.getSuspendNodeIds()));
}
Set<String> suspendNodeIds = newState.getSuspendNodeIds(); EnumSet<ChainStateField> updatedFields = EnumSet.of(
if (suspendNodeIds != null && !suspendNodeIds.isEmpty()) { ChainStateField.STATUS,
// 移除 suspend 状态,方便二次 suspend 时,不带有旧数据 ChainStateField.SUSPEND_NODE_IDS,
updateStateSafely(state -> { 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.setSuspendNodeIds(null);
state.setSuspendForParameters(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); Node node = definition.getNodeById(id);
if (node == null) { if (node == null) {
throw new ChainException("Node not found: " + id); throw new ChainException("Node not found: " + id);
} }
scheduleNode(node, null, TriggerType.RESUME, 0L); scheduleNode(node, null, TriggerType.RESUME, 0L);
} }
} return true;
} }
public void resume() { 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

@@ -42,6 +42,7 @@ public abstract class Node implements Serializable {
protected NodeCondition condition; protected NodeCondition condition;
protected NodeValidator validator; protected NodeValidator validator;
protected NodeJoinMode joinMode = NodeJoinMode.ANY;
// 循环执行相关属性 // 循环执行相关属性
protected boolean loopEnable = false; // 是否启用循环执行 protected boolean loopEnable = false; // 是否启用循环执行
@@ -70,6 +71,10 @@ public abstract class Node implements Serializable {
} }
public void setParentId(String parentId) { public void setParentId(String parentId) {
if (StringUtil.hasText(parentId) && getJoinMode() == NodeJoinMode.ALL) {
throw new IllegalArgumentException(
"joinMode 'all' is not supported for loop child nodes");
}
this.parentId = parentId; this.parentId = parentId;
} }
@@ -121,6 +126,31 @@ public abstract class Node implements Serializable {
this.validator = validator; this.validator = validator;
} }
/**
* 获取节点的直接入边汇聚模式。
*
* <p>旧序列化对象缺少该字段时返回 {@link NodeJoinMode#ANY}。</p>
*
* @return 汇聚模式
*/
public NodeJoinMode getJoinMode() {
return joinMode == null ? NodeJoinMode.ANY : joinMode;
}
/**
* 设置节点的直接入边汇聚模式。
*
* @param joinMode 汇聚模式
*/
public void setJoinMode(NodeJoinMode joinMode) {
NodeJoinMode resolved = joinMode == null ? NodeJoinMode.ANY : joinMode;
if (resolved == NodeJoinMode.ALL && StringUtil.hasText(parentId)) {
throw new IllegalArgumentException(
"joinMode 'all' is not supported for loop child nodes");
}
this.joinMode = resolved;
}
// protected void addOutwardEdge(Edge edge) { // protected void addOutwardEdge(Edge edge) {
// if (this.outwardEdges == null) { // if (this.outwardEdges == null) {
// this.outwardEdges = new ArrayList<>(); // this.outwardEdges = new ArrayList<>();

View File

@@ -0,0 +1,60 @@
/**
* Copyright (c) 2025-2026, Michael Yang 杨福海 (fuhai999@gmail.com).
* <p>
* Licensed under the GNU Lesser General Public License (LGPL) ,Version 3.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
* <p>
* http://www.gnu.org/licenses/lgpl-3.0.txt
* <p>
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package com.easyagents.flow.core.chain;
import java.util.Locale;
/**
* 多入边节点的触发汇聚模式。
*/
public enum NodeJoinMode {
/** 任意一条直接入边到达即可执行。 */
ANY("any"),
/** 全部直接入边到达后才执行。 */
ALL("all");
private final String value;
NodeJoinMode(String value) {
this.value = value;
}
public String getValue() {
return value;
}
/**
* 按工作流 JSON 值解析汇聚模式。
*
* @param value 配置值
* @return 汇聚模式
* @throws IllegalArgumentException 配置为空或不受支持
*/
public static NodeJoinMode ofValue(String value) {
if (value == null || value.trim().isEmpty()) {
throw new IllegalArgumentException("joinMode must be 'any' or 'all'");
}
String normalized = value.trim().toLowerCase(Locale.ROOT);
for (NodeJoinMode mode : values()) {
if (mode.value.equals(normalized)) {
return mode;
}
}
throw new IllegalArgumentException(
"Unsupported joinMode: " + value + "; expected 'any' or 'all'");
}
}

View File

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

@@ -18,6 +18,7 @@ package com.easyagents.flow.core.knowledge;
import java.util.ArrayList; import java.util.ArrayList;
import java.util.List; import java.util.List;
import java.util.Map;
public class KnowledgeManager { public class KnowledgeManager {
@@ -51,4 +52,20 @@ public class KnowledgeManager {
} }
return null; 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; package com.easyagents.flow.core.knowledge;
import com.easyagents.flow.core.util.Maps;
import java.util.List;
import java.util.Map;
public interface KnowledgeProvider { public interface KnowledgeProvider {
Knowledge getKnowledge(Object id); 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; package com.easyagents.flow.core.node;
import com.easyagents.flow.core.chain.Chain; import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainSuspendException; 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.Parameter;
import com.easyagents.flow.core.chain.ParameterOption;
import com.easyagents.flow.core.chain.RefType; import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.ChainStateField; 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 { public class ConfirmNode extends BaseNode {
private static final long serialVersionUID = 1L; 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 String message;
private List<Parameter> confirms; private boolean multiple;
private List<String> options;
public String getMessage() { public String getMessage() {
return message; return message;
@@ -39,115 +73,152 @@ public class ConfirmNode extends BaseNode {
this.message = message; this.message = message;
} }
public List<Parameter> getConfirms() { public boolean isMultiple() {
return confirms; return multiple;
} }
public void setConfirms(List<Parameter> confirms) { public void setMultiple(boolean multiple) {
if (confirms != null) { this.multiple = multiple;
for (Parameter confirm : confirms) {
confirm.setRefType(RefType.INPUT);
confirm.setRequired(true); // 必填,才能正确通过 getParameterValuesOnly 获取参数值
confirm.setName(confirm.getName());
}
}
this.confirms = confirms;
} }
public List<String> getOptions() {
return options;
}
public void setOptions(List<String> options) {
this.options = options;
}
@Override @Override
public Map<String, Object> execute(Chain chain) { public Map<String, Object> execute(Chain chain) {
validateConfiguration();
String outputName = resolveOutputName();
Parameter parameter = buildParameter();
List<Parameter> confirmParameters = new ArrayList<>(); // 确认值只能来自经过 Chain.resumeIfSuspended 校验后创建的恢复触发器。
addConfirmParameter(confirmParameters); // 启动参数与普通节点内存中的同名值均不能绕过人工确认。
if (!isValidatedResumeTrigger(chain)) {
if (confirms != null) { chain.updateStateSafely(state -> {
for (Parameter confirm : confirms) { if (state.getMemory().remove(parameter.getName()) == null) {
Parameter clone = confirm.clone(); return null;
clone.setName(confirm.getName() + "__" + getId());
clone.setRefType(RefType.INPUT);
confirmParameters.add(clone);
} }
return EnumSet.of(ChainStateField.MEMORY);
});
} }
Map<String, Object> values; Map<String, Object> values;
try { try {
values = chain.getExecutionState() values = chain.getExecutionState()
.resolveParameters(this, confirmParameters); .resolveParameters(this, Collections.singletonList(parameter));
// 移除 confirm 参数,方便在其他节点二次确认,或者在 for 循环中第二次获取
chain.updateStateSafely(state -> { chain.updateStateSafely(state -> {
for (Parameter confirmParameter : confirmParameters) { if (!state.getMemory().containsKey(parameter.getName())) {
state.getMemory().remove(confirmParameter.getName()); return null;
} }
state.getMemory().remove(parameter.getName());
return EnumSet.of(ChainStateField.MEMORY); return EnumSet.of(ChainStateField.MEMORY);
}); });
} catch (ChainSuspendException e) { } catch (ChainSuspendException exception) {
chain.updateStateSafely(state -> { chain.updateStateSafely(state -> {
state.setMessage(message); state.setMessage(message);
return EnumSet.of(ChainStateField.MESSAGE); return EnumSet.of(ChainStateField.MESSAGE);
}); });
throw exception;
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);
} }
// 获取参数值,不会触发 ChainSuspendException 错误 return Collections.singletonMap(
Map<String, Object> parameterValues = outputName,
chain.getExecutionState().resolveParameters( values.get(parameter.getName()));
this, }
newParameters,
null,
true);
// 设置 enums方便前端给用户进行选择 /**
for (Parameter confirmParameter : confirmParameters) { * 获取并校验当前节点配置的唯一输出名称。
if (confirmParameter.getEnums() == null) { *
Object enumsObject = parameterValues.get(confirmParameter.getName()); * @return 用户配置的输出名称
confirmParameter.setEnumsObject(enumsObject); */
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; private Parameter buildParameter() {
}
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) {
// “确认 和 取消” 的参数
Parameter parameter = new Parameter(); 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.setRefType(RefType.INPUT);
parameter.setId("confirm");
parameter.setName("confirm__" + getId());
parameter.setRequired(true); parameter.setRequired(true);
List<Object> selectionData = new ArrayList<>();
selectionData.add("yes");
selectionData.add("no");
parameter.setEnums(selectionData);
parameter.setContentType("text"); parameter.setContentType("text");
parameter.setFormType("confirm"); parameter.setFormType(multiple ? "checkbox" : "radio");
parameters.add(parameter); 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.Chain;
import com.easyagents.flow.core.chain.ChainState; 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.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.StringUtil;
import com.easyagents.flow.core.util.TextTemplate; import com.easyagents.flow.core.util.TextTemplate;
import org.slf4j.Logger; import org.slf4j.Logger;
import java.util.Arrays; import java.util.ArrayList;
import java.util.Collections; import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List; import java.util.List;
import java.util.Map; 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 static final Logger logger = org.slf4j.LoggerFactory.getLogger(KnowledgeNode.class);
private Object knowledgeId; private Object knowledgeId;
private List<Object> knowledgeIds = new ArrayList<>();
private String keyword; private String keyword;
private String limit; private String limit;
private String retrievalMode = "HYBRID"; private String retrievalMode = "HYBRID";
@@ -48,6 +49,32 @@ public class KnowledgeNode extends BaseNode {
this.knowledgeId = knowledgeId; 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() { public String getKeyword() {
return keyword; return keyword;
} }
@@ -88,25 +115,44 @@ public class KnowledgeNode extends BaseNode {
if (StringUtil.hasText(realLimitString)) { if (StringUtil.hasText(realLimitString)) {
try { try {
realLimit = Integer.parseInt(realLimitString); realLimit = Integer.parseInt(realLimitString);
} catch (Exception e) { } catch (NumberFormatException exception) {
logger.error(e.toString(), e); throw new IllegalArgumentException(
"知识库节点最终返回条数必须为正整数", exception);
} }
} }
if (realLimit <= 0) {
Knowledge knowledge = KnowledgeManager.getInstance().getKnowledge(knowledgeId); throw new IllegalArgumentException(
"知识库节点最终返回条数必须为正整数");
if (knowledge == null) {
return Collections.emptyMap();
} }
List<Map<String, Object>> result = knowledge.search(realKeyword, realLimit, this, chain); List<Object> resolvedKnowledgeIds = getKnowledgeIds();
return Maps.of("documents", result); 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 @Override
public String toString() { public String toString() {
return "KnowledgeNode{" + return "KnowledgeNode{" +
"knowledgeId=" + knowledgeId + "knowledgeId=" + knowledgeId +
", knowledgeIds=" + knowledgeIds +
", keyword='" + keyword + '\'' + ", keyword='" + keyword + '\'' +
", limit='" + limit + '\'' + ", limit='" + limit + '\'' +
", retrievalMode='" + retrievalMode + '\'' + ", retrievalMode='" + retrievalMode + '\'' +

View File

@@ -21,6 +21,7 @@ import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.DataType; import com.easyagents.flow.core.chain.DataType;
import com.easyagents.flow.core.chain.JsCodeCondition; import com.easyagents.flow.core.chain.JsCodeCondition;
import com.easyagents.flow.core.chain.Node; import com.easyagents.flow.core.chain.Node;
import com.easyagents.flow.core.chain.NodeJoinMode;
import com.easyagents.flow.core.chain.Parameter; import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType; import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.node.BaseNode; import com.easyagents.flow.core.node.BaseNode;
@@ -123,6 +124,10 @@ public abstract class BaseNodeParser<T extends BaseNode> implements NodeParser<T
if (!data.isEmpty()) { if (!data.isEmpty()) {
if (data.containsKey("joinMode")) {
node.setJoinMode(NodeJoinMode.ofValue(data.getString("joinMode")));
}
addParameters(node, data); addParameters(node, data);
addOutputDefs(node, data); addOutputDefs(node, data);

View File

@@ -15,11 +15,12 @@
*/ */
package com.easyagents.flow.core.parser.impl; package com.easyagents.flow.core.parser.impl;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject; import com.alibaba.fastjson.JSONObject;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.node.ConfirmNode; import com.easyagents.flow.core.node.ConfirmNode;
import com.easyagents.flow.core.parser.BaseNodeParser; import com.easyagents.flow.core.parser.BaseNodeParser;
import java.util.ArrayList;
import java.util.List; import java.util.List;
public class ConfirmNodeParser extends BaseNodeParser<ConfirmNode> { 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) { public ConfirmNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) {
ConfirmNode confirmNode = new ConfirmNode(); ConfirmNode confirmNode = new ConfirmNode();
confirmNode.setMessage(data.getString("message")); for (String key : data.keySet()) {
if (!ConfirmNode.SUPPORTED_CONFIGURATION_KEYS.contains(key)) {
List<Parameter> confirms = getParameters(data, "confirms"); throw new IllegalArgumentException(
if (confirms != null && !confirms.isEmpty()) { "用户确认节点包含无效配置字段: " + key);
confirmNode.setConfirms(confirms);
} }
}
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; return confirmNode;
} }

View File

@@ -16,15 +16,41 @@
package com.easyagents.flow.core.parser.impl; package com.easyagents.flow.core.parser.impl;
import com.alibaba.fastjson.JSONObject; import com.alibaba.fastjson.JSONObject;
import com.alibaba.fastjson.JSONArray;
import com.easyagents.flow.core.node.KnowledgeNode; import com.easyagents.flow.core.node.KnowledgeNode;
import com.easyagents.flow.core.parser.BaseNodeParser; import com.easyagents.flow.core.parser.BaseNodeParser;
import java.util.ArrayList;
public class KnowledgeNodeParser extends BaseNodeParser<KnowledgeNode> { public class KnowledgeNodeParser extends BaseNodeParser<KnowledgeNode> {
@Override @Override
public KnowledgeNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) { public KnowledgeNode doParse(JSONObject root, JSONObject data, JSONObject chainJSONObject) {
KnowledgeNode knowledgeNode = new KnowledgeNode(); 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.setKnowledgeId(data.get("knowledgeId"));
}
knowledgeNode.setLimit(data.getString("limit")); knowledgeNode.setLimit(data.getString("limit"));
knowledgeNode.setKeyword(data.getString("keyword")); knowledgeNode.setKeyword(data.getString("keyword"));
knowledgeNode.setRetrievalMode(data.getString("retrievalMode")); knowledgeNode.setRetrievalMode(data.getString("retrievalMode"));

View File

@@ -15,10 +15,15 @@
*/ */
package com.easyagents.flow.core.test; 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.Chain;
import com.easyagents.flow.core.chain.ChainDefinition; import com.easyagents.flow.core.chain.ChainDefinition;
import com.easyagents.flow.core.chain.ChainStatus; 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.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.ChainState;
import com.easyagents.flow.core.chain.event.ChainEndEvent; import com.easyagents.flow.core.chain.event.ChainEndEvent;
import com.easyagents.flow.core.chain.repository.ChainDefinitionSnapshotRepository; 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.BaseNode;
import com.easyagents.flow.core.node.ConfirmNode; import com.easyagents.flow.core.node.ConfirmNode;
import com.easyagents.flow.core.node.StartNode; import com.easyagents.flow.core.node.StartNode;
import com.easyagents.flow.core.parser.ChainParser;
import org.junit.Assert; import org.junit.Assert;
import org.junit.Test; import org.junit.Test;
@@ -61,6 +67,127 @@ import java.util.concurrent.atomic.AtomicReference;
*/ */
public class ChainExecutorConcurrencyTest { 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 入口遇到人工挂起会快速失败,不会无限占用调用线程。 * 验证同步 Tool 入口遇到人工挂起会快速失败,不会无限占用调用线程。
* *
@@ -563,6 +690,12 @@ public class ChainExecutorConcurrencyTest {
start.setId("start"); start.setId("start");
ConfirmNode confirm = new ConfirmNode(); ConfirmNode confirm = new ConfirmNode();
confirm.setId("confirm"); 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(); EndNode end = new EndNode();
end.setId("end"); end.setId("end");
Edge first = new Edge(); Edge first = new Edge();
@@ -581,6 +714,91 @@ public class ChainExecutorConcurrencyTest {
return definition; 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.Chain;
import com.easyagents.flow.core.chain.ChainDefinition; 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.ChainStatus;
import com.easyagents.flow.core.chain.EventManager; 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.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository; import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import org.junit.Assert; import org.junit.Assert;
import org.junit.Test; import org.junit.Test;
import java.util.Arrays;
import java.util.Collections; import java.util.Collections;
import java.util.EnumSet;
import java.util.List;
import java.util.Map; 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} 暂停恢复状态守卫测试。 * {@link Chain} 暂停恢复状态守卫测试。
@@ -72,6 +86,167 @@ public class ChainResumeGuardTest {
.containsKey("unexpected")); .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()); chain.setEventManager(new EventManager());
return chain; 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,407 @@
package com.easyagents.flow.core.test;
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.chain.ChainStatus;
import com.easyagents.flow.core.chain.Edge;
import com.easyagents.flow.core.chain.Node;
import com.easyagents.flow.core.chain.NodeCondition;
import com.easyagents.flow.core.chain.NodeJoinMode;
import com.easyagents.flow.core.chain.NodeState;
import com.easyagents.flow.core.chain.Parameter;
import com.easyagents.flow.core.chain.RefType;
import com.easyagents.flow.core.chain.repository.InMemoryChainStateRepository;
import com.easyagents.flow.core.chain.repository.InMemoryNodeStateRepository;
import com.easyagents.flow.core.chain.runtime.ChainExecutor;
import com.easyagents.flow.core.chain.runtime.InMemoryTriggerStore;
import com.easyagents.flow.core.chain.runtime.TriggerScheduler;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.node.EndNode;
import com.easyagents.flow.core.node.StartNode;
import com.easyagents.flow.core.parser.ChainParser;
import com.easyagents.flow.core.parser.impl.EndNodeParser;
import org.junit.Assert;
import org.junit.Test;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BooleanSupplier;
/**
* 验证普通节点的直接入边汇聚模式。
*/
public class NodeJoinModeTest {
@Test
public void shouldParseJoinModeAndRejectInvalidValues() {
ChainParser parser = ChainParser.builder()
.withDefaultParsers(true)
.build();
Assert.assertEquals(
NodeJoinMode.ANY,
parseEndNode(parser, null, null).getJoinMode());
Assert.assertEquals(
NodeJoinMode.ALL,
parseEndNode(parser, "all", null).getJoinMode());
Assert.assertEquals(
NodeJoinMode.ANY,
parseEndNode(parser, "ANY", null).getJoinMode());
assertInvalidJoinMode(() -> parseEndNode(parser, "first", null));
assertInvalidJoinMode(() -> parseEndNode(parser, "", null));
assertInvalidJoinMode(() -> parseEndNode(parser, "all", "loop"));
ProbeJoinNode node = new ProbeJoinNode(new AtomicInteger());
node.setJoinMode(NodeJoinMode.ALL);
assertInvalidJoinMode(() -> node.setParentId("loop"));
}
@Test
public void shouldWaitForEveryInboundEdgeBeforeExecutingAndCheckingCondition()
throws Exception {
JoinFixture fixture = createJoinFixture(NodeJoinMode.ALL, true);
String instanceId = null;
try {
instanceId = fixture.executor.executeAsync(
fixture.definition.getId(), Collections.emptyMap());
Assert.assertTrue(fixture.branchACompleted.await(2, TimeUnit.SECONDS));
Assert.assertTrue(fixture.branchBStarted.await(2, TimeUnit.SECONDS));
String currentInstanceId = instanceId;
await(() -> hasTriggerEdge(
fixture.nodeStateRepository,
currentInstanceId,
"join",
"a-join"));
Assert.assertEquals(0, fixture.joinExecutions.get());
Assert.assertEquals(0, fixture.conditionChecks.get());
fixture.releaseBranchB.countDown();
ChainState finalState = awaitTerminal(
fixture.chainStateRepository, instanceId);
Assert.assertEquals(1, fixture.joinExecutions.get());
Assert.assertEquals(1, fixture.conditionChecks.get());
Assert.assertEquals("A+B", finalState.getExecuteResult().get("combined"));
Assert.assertEquals(Boolean.TRUE, finalState.getExecuteResult().get("sawBoth"));
} finally {
fixture.releaseBranchB.countDown();
fixture.scheduler.shutdown();
}
}
@Test
public void shouldKeepAnyModeFirstArrivalBehavior() throws Exception {
JoinFixture fixture = createJoinFixture(NodeJoinMode.ANY, false);
try {
String instanceId = fixture.executor.executeAsync(
fixture.definition.getId(), Collections.emptyMap());
Assert.assertTrue(fixture.branchACompleted.await(2, TimeUnit.SECONDS));
Assert.assertTrue(fixture.branchBStarted.await(2, TimeUnit.SECONDS));
await(() -> fixture.joinExecutions.get() > 0);
ChainState finalState = awaitTerminal(
fixture.chainStateRepository, instanceId);
Assert.assertEquals(1, fixture.joinExecutions.get());
Assert.assertEquals(Boolean.FALSE, finalState.getExecuteResult().get("sawBoth"));
} finally {
fixture.releaseBranchB.countDown();
fixture.scheduler.shutdown();
}
}
@Test
public void shouldKeepDefaultAnyRetryAndLoopBehavior() throws Exception {
ScheduledExecutorService schedulerPool =
Executors.newSingleThreadScheduledExecutor();
ExecutorService workerPool = Executors.newFixedThreadPool(3);
TriggerScheduler scheduler = new TriggerScheduler(
new InMemoryTriggerStore(), schedulerPool, workerPool, 1000L);
AtomicInteger executions = new AtomicInteger();
ChainDefinition definition = new ChainDefinition();
definition.setId("join-mode-retry-loop");
StartNode start = new StartNode();
start.setId("start");
RetryLoopNode worker = new RetryLoopNode(executions);
worker.setId("worker");
worker.setRetryEnable(true);
worker.setMaxRetryCount(1);
worker.setRetryIntervalMs(0L);
worker.setLoopEnable(true);
worker.setMaxLoopCount(2);
worker.setLoopIntervalMs(0L);
EndNode end = endNode("worker", "count", "count");
definition.addNode(start);
definition.addNode(worker);
definition.addNode(end);
definition.addEdge(edge("start-worker", "start", "worker"));
definition.addEdge(edge("worker-end", "worker", "end"));
ChainExecutor executor = new ChainExecutor(
ignored -> definition,
new InMemoryChainStateRepository(),
new InMemoryNodeStateRepository(),
scheduler);
try {
Map<String, Object> result = executor.execute(
definition.getId(), Collections.emptyMap(), 5L, TimeUnit.SECONDS);
Assert.assertEquals(NodeJoinMode.ANY, worker.getJoinMode());
Assert.assertEquals(3, executions.get());
Assert.assertEquals(3, result.get("count"));
} finally {
scheduler.shutdown();
}
}
private JoinFixture createJoinFixture(
NodeJoinMode joinMode, boolean withCondition) {
JoinFixture fixture = new JoinFixture();
fixture.schedulerPool = Executors.newScheduledThreadPool(2);
fixture.workerPool = Executors.newFixedThreadPool(4);
fixture.scheduler = new TriggerScheduler(
new InMemoryTriggerStore(),
fixture.schedulerPool,
fixture.workerPool,
1000L);
fixture.chainStateRepository = new InMemoryChainStateRepository();
fixture.nodeStateRepository = new InMemoryNodeStateRepository();
fixture.definition = new ChainDefinition();
fixture.definition.setId("join-mode-" + joinMode.getValue());
StartNode start = new StartNode();
start.setId("start");
BranchNode branchA = new BranchNode(
"A", fixture.branchACompleted, null, null);
branchA.setId("a");
BranchNode branchB = new BranchNode(
"B", null, fixture.branchBStarted, fixture.releaseBranchB);
branchB.setId("b");
ProbeJoinNode join = new ProbeJoinNode(fixture.joinExecutions);
join.setId("join");
join.setJoinMode(joinMode);
if (withCondition) {
join.setCondition(new BothOutputsCondition(fixture.conditionChecks));
}
EndNode end = endNode("join", "combined", "combined");
end.addOutputDef(outputRef("sawBoth", "join.sawBoth"));
fixture.definition.addNode(start);
fixture.definition.addNode(branchA);
fixture.definition.addNode(branchB);
fixture.definition.addNode(join);
fixture.definition.addNode(end);
fixture.definition.addEdge(edge("start-a", "start", "a"));
fixture.definition.addEdge(edge("start-b", "start", "b"));
fixture.definition.addEdge(edge("a-join", "a", "join"));
fixture.definition.addEdge(edge("b-join", "b", "join"));
fixture.definition.addEdge(edge("join-end", "join", "end"));
fixture.executor = new ChainExecutor(
ignored -> fixture.definition,
fixture.chainStateRepository,
fixture.nodeStateRepository,
fixture.scheduler);
return fixture;
}
private Node parseEndNode(
ChainParser parser, String joinMode, String parentId) {
JSONObject data = new JSONObject();
if (joinMode != null) {
data.put("joinMode", joinMode);
}
JSONObject nodeJson = new JSONObject();
nodeJson.put("id", "end");
nodeJson.put("type", "endNode");
nodeJson.put("parentId", parentId);
nodeJson.put("data", data);
return new EndNodeParser().parse(
nodeJson, new JSONObject(), parser);
}
private void assertInvalidJoinMode(Runnable action) {
try {
action.run();
Assert.fail("Expected invalid join mode");
} catch (IllegalArgumentException expected) {
Assert.assertTrue(expected.getMessage().contains("joinMode"));
}
}
private static boolean hasTriggerEdge(
InMemoryNodeStateRepository repository,
String instanceId,
String nodeId,
String edgeId) {
NodeState state = repository.load(instanceId, nodeId);
return state != null && state.getTriggerEdgeIds().contains(edgeId);
}
private static ChainState awaitTerminal(
InMemoryChainStateRepository repository,
String instanceId) throws Exception {
await(() -> {
ChainState state = repository.load(instanceId);
return state != null
&& state.getStatus() != null
&& state.getStatus().isTerminal();
});
ChainState state = repository.load(instanceId);
Assert.assertEquals(ChainStatus.SUCCEEDED, state.getStatus());
return state;
}
private static void await(BooleanSupplier condition) throws Exception {
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(3L);
while (!condition.getAsBoolean() && System.nanoTime() < deadline) {
Thread.sleep(10L);
}
Assert.assertTrue("condition was not met before timeout", condition.getAsBoolean());
}
private static EndNode endNode(
String sourceNodeId, String sourceName, String outputName) {
EndNode end = new EndNode();
end.setId("end");
end.addOutputDef(outputRef(
outputName, sourceNodeId + "." + sourceName));
return end;
}
private static Parameter outputRef(String name, String ref) {
Parameter parameter = new Parameter();
parameter.setName(name);
parameter.setRef(ref);
parameter.setRefType(RefType.REF);
return parameter;
}
private static Edge edge(String id, String source, String target) {
Edge edge = new Edge();
edge.setId(id);
edge.setSource(source);
edge.setTarget(target);
return edge;
}
private static final class BranchNode extends BaseNode {
private final String value;
private final CountDownLatch completed;
private final CountDownLatch started;
private final CountDownLatch release;
private BranchNode(
String value,
CountDownLatch completed,
CountDownLatch started,
CountDownLatch release) {
this.value = value;
this.completed = completed;
this.started = started;
this.release = release;
}
@Override
public Map<String, Object> execute(Chain chain) {
if (started != null) {
started.countDown();
}
if (release != null) {
try {
if (!release.await(3L, TimeUnit.SECONDS)) {
throw new IllegalStateException("branch release timed out");
}
} catch (InterruptedException error) {
Thread.currentThread().interrupt();
throw new IllegalStateException("branch interrupted", error);
}
}
if (completed != null) {
completed.countDown();
}
return Collections.singletonMap("value", value);
}
}
private static final class ProbeJoinNode extends BaseNode {
private final AtomicInteger executions;
private ProbeJoinNode(AtomicInteger executions) {
this.executions = executions;
}
@Override
public Map<String, Object> execute(Chain chain) {
executions.incrementAndGet();
Object a = chain.getExecutionState().getMemory().get("a.value");
Object b = chain.getExecutionState().getMemory().get("b.value");
Map<String, Object> result = new HashMap<>();
result.put("combined", String.valueOf(a) + "+" + String.valueOf(b));
result.put("sawBoth", a != null && b != null);
return result;
}
}
private static final class BothOutputsCondition implements NodeCondition {
private static final long serialVersionUID = 1L;
private final AtomicInteger checks;
private BothOutputsCondition(AtomicInteger checks) {
this.checks = checks;
}
@Override
public boolean check(
Chain chain,
NodeState context,
Map<String, Object> executeResult) {
checks.incrementAndGet();
Map<String, Object> memory = chain.getExecutionState().getMemory();
return memory.containsKey("a.value") && memory.containsKey("b.value");
}
}
private static final class RetryLoopNode extends BaseNode {
private final AtomicInteger executions;
private RetryLoopNode(AtomicInteger executions) {
this.executions = executions;
}
@Override
public Map<String, Object> execute(Chain chain) {
int count = executions.incrementAndGet();
if (count == 1) {
throw new IllegalStateException("retry once");
}
return Collections.singletonMap("count", count);
}
}
private static final class JoinFixture {
private final CountDownLatch branchACompleted = new CountDownLatch(1);
private final CountDownLatch branchBStarted = new CountDownLatch(1);
private final CountDownLatch releaseBranchB = new CountDownLatch(1);
private final AtomicInteger joinExecutions = new AtomicInteger();
private final AtomicInteger conditionChecks = new AtomicInteger();
private ScheduledExecutorService schedulerPool;
private ExecutorService workerPool;
private TriggerScheduler scheduler;
private InMemoryChainStateRepository chainStateRepository;
private InMemoryNodeStateRepository nodeStateRepository;
private ChainDefinition definition;
private ChainExecutor executor;
}
}

View File

@@ -2,14 +2,26 @@ package com.easyagents.rag.retrieval;
import com.easyagents.core.document.Document; import com.easyagents.core.document.Document;
import java.util.ArrayList;
import java.util.List; import java.util.List;
/**
* 将不同检索路径的原始分数转换为统一的零到一最终相关度。
*/
public final class RagScoreNormalizer { public final class RagScoreNormalizer {
/**
* 禁止实例化工具类。
*/
private RagScoreNormalizer() { private RagScoreNormalizer() {
} }
/**
* 按检索模式归一化文档最终分数。
*
* @param documents 待归一化文档
* @param retrievalMode 检索模式
* @param reranked 是否已经过重排模型
*/
public static void normalize(List<Document> documents, RetrievalMode retrievalMode, boolean reranked) { public static void normalize(List<Document> documents, RetrievalMode retrievalMode, boolean reranked) {
if (documents == null || documents.isEmpty()) { if (documents == null || documents.isEmpty()) {
return; return;
@@ -49,41 +61,29 @@ public final class RagScoreNormalizer {
} }
} }
/**
* 保留重排模型返回的绝对相关度,并将异常范围限制到零到一。
*
* @param documents 待归一化的文档
*/
private static void normalizeRerankScores(List<Document> documents) { private static void normalizeRerankScores(List<Document> documents) {
List<Double> rawScores = new ArrayList<Double>(documents.size());
boolean allPresent = true;
Double min = null;
Double max = null;
for (Document document : documents) { for (Document document : documents) {
Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.RERANK_SCORE, document == null ? null : document.getScore()); Double rawScore = readRawScore(document, RagRetrievalMetadataKeys.RERANK_SCORE, document == null ? null : document.getScore());
rawScores.add(rawScore); if (document != null) {
if (rawScore == null) { // Rerank 适配器返回绝对相关度,按查询结果集再次缩放会把低相关第一名错误抬高到 1。
allPresent = false; document.setScore(clamp01(rawScore));
continue;
} }
min = min == null ? rawScore : Math.min(min, rawScore);
max = max == null ? rawScore : Math.max(max, rawScore);
}
if (allPresent && min != null && max != null && Double.compare(max, min) != 0) {
for (int i = 0; i < documents.size(); i++) {
Double rawScore = rawScores.get(i);
documents.get(i).setScore(clamp01((rawScore - min) / (max - min)));
}
return;
}
if (documents.size() == 1) {
documents.get(0).setScore(1D);
return;
}
int size = documents.size();
for (int i = 0; i < size; i++) {
documents.get(i).setScore(clamp01(1D - ((double) i / (double) (size - 1))));
} }
} }
/**
* 优先读取指定元数据中的原始分数。
*
* @param document 文档
* @param metadataKey 分数元数据键
* @param fallback 元数据不可用时的回退分数
* @return 原始分数;文档为空时返回 null
*/
private static Double readRawScore(Document document, String metadataKey, Double fallback) { private static Double readRawScore(Document document, String metadataKey, Double fallback) {
if (document == null) { if (document == null) {
return null; return null;
@@ -102,6 +102,12 @@ public final class RagScoreNormalizer {
return fallback; return fallback;
} }
/**
* 将可空分数限制到零到一范围。
*
* @param value 原始分数
* @return 有效最终分数
*/
private static double clamp01(Double value) { private static double clamp01(Double value) {
if (value == null || value.isNaN() || value.isInfinite()) { if (value == null || value.isNaN() || value.isInfinite()) {
return 0D; return 0D;

View File

@@ -7,8 +7,14 @@ import org.junit.Test;
import java.util.Arrays; import java.util.Arrays;
import java.util.List; import java.util.List;
/**
* {@link RagScoreNormalizer} 回归测试。
*/
public class RagScoreNormalizerTest { public class RagScoreNormalizerTest {
/**
* 验证关键词分数按有界函数归一化。
*/
@Test @Test
public void shouldNormalizeKeywordScoresToZeroAndOneRange() { public void shouldNormalizeKeywordScoresToZeroAndOneRange() {
Document first = document(1, 9D, RagRetrievalMetadataKeys.KEYWORD_SCORE); Document first = document(1, 9D, RagRetrievalMetadataKeys.KEYWORD_SCORE);
@@ -20,6 +26,9 @@ public class RagScoreNormalizerTest {
Assert.assertEquals(0D, second.getScore(), 0.0001D); Assert.assertEquals(0D, second.getScore(), 0.0001D);
} }
/**
* 验证混合检索 RRF 分数按理论上界归一化。
*/
@Test @Test
public void shouldNormalizeHybridFusionScoreByRrfUpperBound() { public void shouldNormalizeHybridFusionScoreByRrfUpperBound() {
Document document = document(1, 2D / (RrfFusionStrategy.DEFAULT_RRF_K + 1D), RagRetrievalMetadataKeys.FUSION_SCORE); Document document = document(1, 2D / (RrfFusionStrategy.DEFAULT_RRF_K + 1D), RagRetrievalMetadataKeys.FUSION_SCORE);
@@ -29,36 +38,60 @@ public class RagScoreNormalizerTest {
Assert.assertEquals(1D, document.getScore(), 0.0001D); Assert.assertEquals(1D, document.getScore(), 0.0001D);
} }
/**
* 验证重排模型返回的绝对相关度保持不变。
*/
@Test @Test
public void shouldNormalizeRerankScoresByMinMax() { public void shouldPreserveRerankRelevanceScores() {
List<Document> documents = Arrays.asList( List<Document> documents = Arrays.asList(
document(1, 10D, RagRetrievalMetadataKeys.RERANK_SCORE), document(1, 0.1D, RagRetrievalMetadataKeys.RERANK_SCORE),
document(2, 20D, RagRetrievalMetadataKeys.RERANK_SCORE), document(2, 0.5D, RagRetrievalMetadataKeys.RERANK_SCORE),
document(3, 30D, RagRetrievalMetadataKeys.RERANK_SCORE) document(3, 0.9D, RagRetrievalMetadataKeys.RERANK_SCORE)
);
RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);
Assert.assertEquals(0.1D, documents.get(0).getScore(), 0.0001D);
Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D);
Assert.assertEquals(0.9D, documents.get(2).getScore(), 0.0001D);
}
/**
* 验证单条低分重排结果不会被抬高为高相关结果。
*/
@Test
public void shouldKeepSingleLowRerankScoreLow() {
Document document = document(1, 0.1D, RagRetrievalMetadataKeys.RERANK_SCORE);
RagScoreNormalizer.normalize(Arrays.asList(document), RetrievalMode.HYBRID, true);
Assert.assertEquals(0.1D, document.getScore(), 0.0001D);
}
/**
* 验证越界重排分数会被限制到零到一范围。
*/
@Test
public void shouldClampRerankScoresToZeroAndOneRange() {
List<Document> documents = Arrays.asList(
document(1, -0.2D, RagRetrievalMetadataKeys.RERANK_SCORE),
document(2, 1.2D, RagRetrievalMetadataKeys.RERANK_SCORE)
); );
RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true); RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);
Assert.assertEquals(0D, documents.get(0).getScore(), 0.0001D); Assert.assertEquals(0D, documents.get(0).getScore(), 0.0001D);
Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D); Assert.assertEquals(1D, documents.get(1).getScore(), 0.0001D);
Assert.assertEquals(1D, documents.get(2).getScore(), 0.0001D);
}
@Test
public void shouldFallbackToRankBasedNormalizationWhenRerankScoresAreEqual() {
List<Document> documents = Arrays.asList(
document(1, 5D, RagRetrievalMetadataKeys.RERANK_SCORE),
document(2, 5D, RagRetrievalMetadataKeys.RERANK_SCORE),
document(3, 5D, RagRetrievalMetadataKeys.RERANK_SCORE)
);
RagScoreNormalizer.normalize(documents, RetrievalMode.HYBRID, true);
Assert.assertEquals(1D, documents.get(0).getScore(), 0.0001D);
Assert.assertEquals(0.5D, documents.get(1).getScore(), 0.0001D);
Assert.assertEquals(0D, documents.get(2).getScore(), 0.0001D);
} }
/**
* 创建携带原始分数元数据的测试文档。
*
* @param id 文档 ID
* @param score 原始分数
* @param metadataKey 分数元数据键
* @return 测试文档
*/
private Document document(Object id, Double score, String metadataKey) { private Document document(Object id, Double score, String metadataKey) {
Document document = new Document(); Document document = new Document();
document.setId(id); document.setId(id);

View File

@@ -0,0 +1,53 @@
package com.easyagents.scheduler;
import java.time.Duration;
/**
* 请求调度提供方持久化重试同一次逻辑触发的异常。
*
* <p>仅用于尚未产生业务副作用、且调用方必须保证至少一次登记的短 Handler。
* {@code maxRefires=0} 表示不限制持久化重试次数Provider 正常情况下应释放当前执行线程,
* 并通过持久化的延迟触发保留原 fire 上下文。仅当 JobStore 无法写入持久重试时,
* Provider 可按相同退避暂时占用当前线程并原地重试,以避免正常完成造成 fire 丢失。</p>
*/
public class ScheduleRefireException extends RuntimeException {
private static final int DEFAULT_MAX_REFIRES = 0;
private static final Duration DEFAULT_BASE_DELAY = Duration.ofMillis(250);
private final int maxRefires;
private final Duration baseDelay;
public ScheduleRefireException(String message, Throwable cause) {
this(message, cause, DEFAULT_MAX_REFIRES, DEFAULT_BASE_DELAY);
}
public ScheduleRefireException(String message, Throwable cause,
int maxRefires, Duration baseDelay) {
super(requireMessage(message), cause);
if (maxRefires < 0) throw new IllegalArgumentException("maxRefires must not be negative");
if (baseDelay == null || baseDelay.isZero() || baseDelay.isNegative()) {
throw new IllegalArgumentException("baseDelay must be positive");
}
this.maxRefires = maxRefires;
this.baseDelay = baseDelay;
}
/**
* @return 最大重新执行次数0 表示不限制
*/
public int maxRefires() {
return maxRefires;
}
public Duration delayFor(int refireNumber) {
int shift = Math.min(4, Math.max(0, refireNumber - 1));
return baseDelay.multipliedBy(1L << shift);
}
private static String requireMessage(String message) {
if (message == null || message.isBlank()) {
throw new IllegalArgumentException("message must not be blank");
}
return message;
}
}

View File

@@ -29,6 +29,9 @@ abstract class AbstractDispatchJob implements InterruptableJob {
throw new JobExecutionException("easy-agents scheduler runtime is not available"); throw new JobExecutionException("easy-agents scheduler runtime is not available");
} }
quartzRuntime.execute(context); quartzRuntime.execute(context);
} catch (JobExecutionException exception) {
// JobExecutionException 继承 SchedulerException必须保留 refire 等控制语义。
throw exception;
} catch (SchedulerException exception) { } catch (SchedulerException exception) {
throw new JobExecutionException("failed to access scheduler runtime", exception, false); throw new JobExecutionException("failed to access scheduler runtime", exception, false);
} finally { } finally {

View File

@@ -7,9 +7,16 @@ import com.easyagents.scheduler.ScheduleExecutionListener;
import com.easyagents.scheduler.ScheduleFireContext; import com.easyagents.scheduler.ScheduleFireContext;
import com.easyagents.scheduler.ScheduleHandler; import com.easyagents.scheduler.ScheduleHandler;
import com.easyagents.scheduler.ScheduleId; import com.easyagents.scheduler.ScheduleId;
import com.easyagents.scheduler.ScheduleRefireException;
import org.quartz.JobDataMap; import org.quartz.JobDataMap;
import org.quartz.JobExecutionContext; import org.quartz.JobExecutionContext;
import org.quartz.JobExecutionException; import org.quartz.JobExecutionException;
import org.quartz.ObjectAlreadyExistsException;
import org.quartz.SchedulerException;
import org.quartz.SimpleScheduleBuilder;
import org.quartz.Trigger;
import org.quartz.TriggerBuilder;
import org.quartz.TriggerKey;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
@@ -19,6 +26,10 @@ import java.util.Collection;
import java.util.LinkedHashMap; import java.util.LinkedHashMap;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.nio.charset.StandardCharsets;
import java.util.Date;
import java.util.UUID;
import java.util.concurrent.locks.LockSupport;
/** /**
* 当前 Scheduler 节点的 Handler 和监听器运行时。 * 当前 Scheduler 节点的 Handler 和监听器运行时。
@@ -55,15 +66,27 @@ final class QuartzRuntime {
void execute(JobExecutionContext quartzContext) throws JobExecutionException { void execute(JobExecutionContext quartzContext) throws JobExecutionException {
JobDataMap data = quartzContext.getMergedJobDataMap(); JobDataMap data = quartzContext.getMergedJobDataMap();
ScheduleDefinition definition = QuartzScheduleMapper.toDefinition(data); ScheduleDefinition definition = QuartzScheduleMapper.toDefinition(data);
Instant actualFireTime = toInstant(quartzContext.getFireTime(), Instant.now()); Instant observedActualFireTime = toInstant(quartzContext.getFireTime(), Instant.now());
Instant scheduledFireTime = instantValue(
data,
QuartzScheduleMapper.KEY_RETRY_SCHEDULED_FIRE_TIME,
toInstant(quartzContext.getScheduledFireTime(), observedActualFireTime)
);
Instant actualFireTime = instantValue(
data,
QuartzScheduleMapper.KEY_RETRY_ACTUAL_FIRE_TIME,
observedActualFireTime
);
ScheduleFireContext context = new ScheduleFireContext( ScheduleFireContext context = new ScheduleFireContext(
definition.id(), definition.id(),
definition.handlerCode(), definition.handlerCode(),
toInstant(quartzContext.getScheduledFireTime(), actualFireTime), scheduledFireTime,
actualFireTime, actualFireTime,
quartzContext.getFireInstanceId(), stringValue(data, QuartzScheduleMapper.KEY_RETRY_FIRE_INSTANCE_ID,
quartzContext.getFireInstanceId()),
stringValue(data, QuartzScheduleMapper.KEY_INVOCATION), stringValue(data, QuartzScheduleMapper.KEY_INVOCATION),
quartzContext.isRecovering(), booleanValue(data, QuartzScheduleMapper.KEY_RETRY_RECOVERING,
quartzContext.isRecovering()),
definition.parameters() definition.parameters()
); );
ScheduleHandler handler = handlers.get(definition.handlerCode()); ScheduleHandler handler = handlers.get(definition.handlerCode());
@@ -82,10 +105,80 @@ final class QuartzRuntime {
notifySucceeded(context, elapsed(startedAt)); notifySucceeded(context, elapsed(startedAt));
} catch (Exception exception) { } catch (Exception exception) {
notifyFailed(context, elapsed(startedAt), exception); notifyFailed(context, elapsed(startedAt), exception);
if (exception instanceof ScheduleRefireException retry) {
int attempt = Math.max(
intValue(data, QuartzScheduleMapper.KEY_RETRY_ATTEMPT, 0),
quartzContext.getRefireCount()
);
if (retry.maxRefires() == 0 || attempt < retry.maxRefires()) {
try {
persistRetry(quartzContext, context, retry, attempt + 1);
return;
} catch (SchedulerException retryFailure) {
exception.addSuppressed(retryFailure);
log.error("Failed to persist schedule handler retry: scheduleId={}, attempt={}",
context.scheduleId(), attempt + 1, retryFailure);
if (canRefireInPlace(quartzContext, retry, attempt + 1)) {
// JobStore 暂时不可写时,正常完成原 fire 会造成登记丢失。
// 仅在这条降级路径短时占用当前 worker一旦持久重试落库即释放。
throw new JobExecutionException(exception, true);
}
}
} else {
log.error("Schedule handler retry limit exhausted: scheduleId={}, retries={}",
context.scheduleId(), retry.maxRefires(), exception);
}
}
// 不在 Quartz worker 内无限 refire。持久化重试失败时保留原异常
// 由 requestRecovery 和集群故障恢复处理未完成的 fired trigger。
throw new JobExecutionException(exception, false); throw new JobExecutionException(exception, false);
} }
} }
private static boolean canRefireInPlace(
JobExecutionContext context,
ScheduleRefireException retry,
int attempt
) {
LockSupport.parkNanos(retry.delayFor(attempt).toNanos());
if (Thread.currentThread().isInterrupted()) return false;
try {
return !context.getScheduler().isShutdown();
} catch (SchedulerException exception) {
return false;
}
}
private static void persistRetry(
JobExecutionContext quartzContext,
ScheduleFireContext context,
ScheduleRefireException retry,
int attempt
) throws SchedulerException {
String source = context.scheduleId() + "|" + context.fireInstanceId() + "|"
+ String.valueOf(context.invocationId()) + "|" + attempt;
String retryName = "retry-" + UUID.nameUUIDFromBytes(
source.getBytes(StandardCharsets.UTF_8));
TriggerKey retryKey = new TriggerKey(
retryName,
QuartzScheduleMapper.GROUP_PREFIX + "retry." + context.scheduleId().namespace()
);
Trigger trigger = TriggerBuilder.newTrigger()
.withIdentity(retryKey)
.forJob(quartzContext.getJobDetail().getKey())
.usingJobData(QuartzScheduleMapper.retryData(context, attempt))
.startAt(Date.from(Instant.now().plus(retry.delayFor(attempt))))
.withSchedule(SimpleScheduleBuilder.simpleSchedule()
.withRepeatCount(0)
.withMisfireHandlingInstructionFireNow())
.build();
try {
quartzContext.getScheduler().scheduleJob(trigger);
} catch (ObjectAlreadyExistsException ignored) {
// 同一原始 fire/attempt 的确定性 key 已落库,即视为持久化成功。
}
}
/** /**
* 向监听器发布 Misfire 事件。 * 向监听器发布 Misfire 事件。
* *
@@ -175,4 +268,24 @@ final class QuartzRuntime {
Object value = data.get(key); Object value = data.get(key);
return value == null ? null : value.toString(); return value == null ? null : value.toString();
} }
private static String stringValue(JobDataMap data, String key, String fallback) {
String value = stringValue(data, key);
return value == null ? fallback : value;
}
private static Instant instantValue(JobDataMap data, String key, Instant fallback) {
String value = stringValue(data, key);
return value == null ? fallback : Instant.ofEpochMilli(Long.parseLong(value));
}
private static boolean booleanValue(JobDataMap data, String key, boolean fallback) {
String value = stringValue(data, key);
return value == null ? fallback : Boolean.parseBoolean(value);
}
private static int intValue(JobDataMap data, String key, int fallback) {
String value = stringValue(data, key);
return value == null ? fallback : Integer.parseInt(value);
}
} }

View File

@@ -7,6 +7,7 @@ import com.easyagents.scheduler.OnceSchedulePlan;
import com.easyagents.scheduler.ScheduleDefinition; import com.easyagents.scheduler.ScheduleDefinition;
import com.easyagents.scheduler.ScheduleErrorCode; import com.easyagents.scheduler.ScheduleErrorCode;
import com.easyagents.scheduler.ScheduleException; import com.easyagents.scheduler.ScheduleException;
import com.easyagents.scheduler.ScheduleFireContext;
import com.easyagents.scheduler.ScheduleId; import com.easyagents.scheduler.ScheduleId;
import com.easyagents.scheduler.SchedulePlan; import com.easyagents.scheduler.SchedulePlan;
import org.quartz.CronScheduleBuilder; import org.quartz.CronScheduleBuilder;
@@ -47,6 +48,11 @@ final class QuartzScheduleMapper {
static final String KEY_DESCRIPTION = "ea.description"; static final String KEY_DESCRIPTION = "ea.description";
static final String KEY_INVOCATION = "ea.invocationId"; static final String KEY_INVOCATION = "ea.invocationId";
static final String KEY_IMMEDIATE_PARAMETER_SNAPSHOT = "ea.immediateParameterSnapshot"; static final String KEY_IMMEDIATE_PARAMETER_SNAPSHOT = "ea.immediateParameterSnapshot";
static final String KEY_RETRY_ATTEMPT = "ea.retry.attempt";
static final String KEY_RETRY_SCHEDULED_FIRE_TIME = "ea.retry.scheduledFireTime";
static final String KEY_RETRY_ACTUAL_FIRE_TIME = "ea.retry.actualFireTime";
static final String KEY_RETRY_FIRE_INSTANCE_ID = "ea.retry.fireInstanceId";
static final String KEY_RETRY_RECOVERING = "ea.retry.recovering";
static final String PARAMETER_PREFIX = "ea.parameter."; static final String PARAMETER_PREFIX = "ea.parameter.";
static final String IMMEDIATE_PARAMETER_PREFIX = "ea.immediateParameter."; static final String IMMEDIATE_PARAMETER_PREFIX = "ea.immediateParameter.";
static final String GROUP_PREFIX = "ea.scheduler."; static final String GROUP_PREFIX = "ea.scheduler.";
@@ -153,6 +159,30 @@ final class QuartzScheduleMapper {
return data; return data;
} }
/**
* 固化持久化重试所需的原始 fire 上下文。
*/
static JobDataMap retryData(ScheduleFireContext context, int attempt) {
JobDataMap data = identityData(context.scheduleId());
// Trigger 数据覆盖 JobDetail任务定义被 replace 后,既有 fire 的重试仍须
// 派发给首次触发时的 Handler而不是意外切换到新 Handler。
data.put(KEY_HANDLER, context.handlerCode());
data.put(KEY_RETRY_ATTEMPT, Integer.toString(attempt));
data.put(KEY_RETRY_SCHEDULED_FIRE_TIME,
Long.toString(context.scheduledFireTime().toEpochMilli()));
data.put(KEY_RETRY_ACTUAL_FIRE_TIME,
Long.toString(context.actualFireTime().toEpochMilli()));
data.put(KEY_RETRY_FIRE_INSTANCE_ID, context.fireInstanceId());
data.put(KEY_RETRY_RECOVERING, Boolean.toString(context.recovering()));
if (context.invocationId() != null) {
data.put(KEY_INVOCATION, context.invocationId());
}
// 重试必须沿用首次 fire 的参数快照,不能读取期间被替换的新定义参数。
data.put(KEY_IMMEDIATE_PARAMETER_SNAPSHOT, Boolean.TRUE.toString());
putParameters(data, context.parameters(), IMMEDIATE_PARAMETER_PREFIX);
return data;
}
/** /**
* 从持久 JobData 恢复公共定义。 * 从持久 JobData 恢复公共定义。
* *

View File

@@ -335,7 +335,8 @@ public final class QuartzScheduleService implements ScheduleService, AutoCloseab
} }
try { try {
if (!scheduler.isShutdown()) { if (!scheduler.isShutdown()) {
// false 保证 close 本身有界Quartz 会向内部 InterruptableJob 发送中断。 // Factory 启用 interruptJobsOnShutdownQuartz 会先中断内部
// InterruptableJobfalse 只表示不再无界等待忽略中断的 Handler。
scheduler.shutdown(false); scheduler.shutdown(false);
} }
} catch (SchedulerException exception) { } catch (SchedulerException exception) {

View File

@@ -10,6 +10,8 @@ package com.easyagents.scheduler.quartz;
* @param clustered 是否启用 JDBC 集群 * @param clustered 是否启用 JDBC 集群
* @param threadCount Quartz Worker 线程数 * @param threadCount Quartz Worker 线程数
* @param threadPriority Quartz Worker 线程优先级 * @param threadPriority Quartz Worker 线程优先级
* @param batchTriggerAcquisitionMaxCount 单次批量获取 Trigger 的最大数量
* @param batchTriggerAcquisitionFireAheadTimeWindowMillis 可提前纳入批量的时间窗口,单位毫秒
* @param clusterCheckinIntervalMillis 集群心跳间隔,单位毫秒 * @param clusterCheckinIntervalMillis 集群心跳间隔,单位毫秒
* @param misfireThresholdMillis Misfire 判定阈值,单位毫秒 * @param misfireThresholdMillis Misfire 判定阈值,单位毫秒
* @param waitForJobsToCompleteOnShutdown 关闭时是否等待运行中任务完成 * @param waitForJobsToCompleteOnShutdown 关闭时是否等待运行中任务完成
@@ -24,6 +26,8 @@ public record QuartzSchedulerConfig(
boolean clustered, boolean clustered,
int threadCount, int threadCount,
int threadPriority, int threadPriority,
int batchTriggerAcquisitionMaxCount,
long batchTriggerAcquisitionFireAheadTimeWindowMillis,
long clusterCheckinIntervalMillis, long clusterCheckinIntervalMillis,
long misfireThresholdMillis, long misfireThresholdMillis,
boolean waitForJobsToCompleteOnShutdown, boolean waitForJobsToCompleteOnShutdown,
@@ -55,6 +59,18 @@ public record QuartzSchedulerConfig(
if (threadPriority < Thread.MIN_PRIORITY || threadPriority > Thread.MAX_PRIORITY) { if (threadPriority < Thread.MIN_PRIORITY || threadPriority > Thread.MAX_PRIORITY) {
throw new IllegalArgumentException("threadPriority must be between 1 and 10"); throw new IllegalArgumentException("threadPriority must be between 1 and 10");
} }
if (batchTriggerAcquisitionMaxCount < 1
|| batchTriggerAcquisitionMaxCount > threadCount) {
throw new IllegalArgumentException(
"batchTriggerAcquisitionMaxCount must be between 1 and threadCount"
);
}
if (batchTriggerAcquisitionFireAheadTimeWindowMillis < 0
|| batchTriggerAcquisitionFireAheadTimeWindowMillis > 60_000L) {
throw new IllegalArgumentException(
"batchTriggerAcquisitionFireAheadTimeWindowMillis must be between 0 and 60000"
);
}
if (clusterCheckinIntervalMillis < 1000) { if (clusterCheckinIntervalMillis < 1000) {
throw new IllegalArgumentException( throw new IllegalArgumentException(
"clusterCheckinIntervalMillis must be at least 1000" "clusterCheckinIntervalMillis must be at least 1000"
@@ -85,6 +101,8 @@ public record QuartzSchedulerConfig(
true, true,
8, 8,
Thread.NORM_PRIORITY, Thread.NORM_PRIORITY,
1,
0L,
15_000L, 15_000L,
60_000L, 60_000L,
true, true,

View File

@@ -154,7 +154,7 @@ public final class QuartzSchedulerFactory {
* @param dataSourceName Quartz 内部 DataSource 名称 * @param dataSourceName Quartz 内部 DataSource 名称
* @return Quartz 属性 * @return Quartz 属性
*/ */
private static Properties properties( static Properties properties(
QuartzSchedulerConfig config, QuartzSchedulerConfig config,
String dataSourceName String dataSourceName
) { ) {
@@ -173,6 +173,14 @@ public final class QuartzSchedulerFactory {
"org.quartz.threadPool.threadPriority", "org.quartz.threadPool.threadPriority",
Integer.toString(config.threadPriority()) Integer.toString(config.threadPriority())
); );
properties.setProperty(
"org.quartz.scheduler.batchTriggerAcquisitionMaxCount",
Integer.toString(config.batchTriggerAcquisitionMaxCount())
);
properties.setProperty(
"org.quartz.scheduler.batchTriggerAcquisitionFireAheadTimeWindow",
Long.toString(config.batchTriggerAcquisitionFireAheadTimeWindowMillis())
);
properties.setProperty( properties.setProperty(
"org.quartz.jobStore.class", "org.quartz.jobStore.class",
"org.quartz.impl.jdbcjobstore.JobStoreTX" "org.quartz.impl.jdbcjobstore.JobStoreTX"
@@ -182,6 +190,10 @@ public final class QuartzSchedulerFactory {
config.driverDelegateClass() config.driverDelegateClass()
); );
properties.setProperty("org.quartz.jobStore.useProperties", "true"); properties.setProperty("org.quartz.jobStore.useProperties", "true");
properties.setProperty(
"org.quartz.jobStore.acquireTriggersWithinLock",
Boolean.toString(config.batchTriggerAcquisitionMaxCount() > 1)
);
properties.setProperty("org.quartz.jobStore.dataSource", dataSourceName); properties.setProperty("org.quartz.jobStore.dataSource", dataSourceName);
properties.setProperty("org.quartz.jobStore.tablePrefix", config.tablePrefix()); properties.setProperty("org.quartz.jobStore.tablePrefix", config.tablePrefix());
properties.setProperty( properties.setProperty(

View File

@@ -9,6 +9,7 @@ import com.easyagents.scheduler.ScheduleException;
import com.easyagents.scheduler.ScheduleFireContext; import com.easyagents.scheduler.ScheduleFireContext;
import com.easyagents.scheduler.ScheduleHandler; import com.easyagents.scheduler.ScheduleHandler;
import com.easyagents.scheduler.ScheduleId; import com.easyagents.scheduler.ScheduleId;
import com.easyagents.scheduler.ScheduleRefireException;
import org.h2.jdbcx.JdbcDataSource; import org.h2.jdbcx.JdbcDataSource;
import org.h2.tools.RunScript; import org.h2.tools.RunScript;
import org.junit.Test; import org.junit.Test;
@@ -20,9 +21,11 @@ import java.sql.Connection;
import java.sql.ResultSet; import java.sql.ResultSet;
import java.sql.Statement; import java.sql.Statement;
import java.time.Instant; import java.time.Instant;
import java.time.Duration;
import java.util.List; import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch; import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
@@ -218,6 +221,83 @@ public class QuartzJdbcIntegrationTest {
} }
} }
/**
* 验证延迟重试先落入 JDBC JobStore关闭并重建节点后仍使用原 fire 上下文执行。
*/
@Test
public void shouldRecoverPersistedRetryAfterSchedulerRestart() throws Exception {
JdbcDataSource dataSource = dataSource();
executeSchema(dataSource);
QuartzSchedulerConfig config = jdbcConfig();
CountDownLatch firstFailed = new CountDownLatch(1);
CountDownLatch retryCompleted = new CountDownLatch(1);
CopyOnWriteArrayList<ScheduleFireContext> contexts = new CopyOnWriteArrayList<>();
ScheduleHandler retryingHandler = new ScheduleHandler() {
@Override
public String code() {
return "persistent-retry-handler";
}
@Override
public void execute(ScheduleFireContext context) {
contexts.add(context);
if (contexts.size() == 1) {
firstFailed.countDown();
throw new ScheduleRefireException("temporary database failure",
new IllegalStateException("unavailable"), 2,
Duration.ofSeconds(2));
}
retryCompleted.countDown();
}
};
ScheduleId scheduleId = new ScheduleId("jdbc", "persistent-retry");
QuartzScheduleService first = QuartzSchedulerFactory.createJdbc(
dataSource, config, List.of(retryingHandler), List.of());
first.start();
first.create(new ScheduleDefinition(
scheduleId,
retryingHandler.code(),
new OnceSchedulePlan(Instant.parse("2099-01-01T00:00:00Z")),
MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW,
true,
Map.of("snapshot", "original"),
"persistent retry"
));
first.triggerNow(scheduleId, "persistent-invocation", Map.of());
assertTrue(firstFailed.await(5, TimeUnit.SECONDS));
awaitRetryTrigger(dataSource);
first.close();
QuartzScheduleService restarted = QuartzSchedulerFactory.createJdbc(
dataSource, config, List.of(retryingHandler), List.of());
try {
restarted.start();
assertTrue("persisted retry did not execute after restart",
retryCompleted.await(8, TimeUnit.SECONDS));
assertEquals(contexts.get(0).scheduledFireTime(), contexts.get(1).scheduledFireTime());
assertEquals(contexts.get(0).actualFireTime(), contexts.get(1).actualFireTime());
assertEquals(contexts.get(0).fireInstanceId(), contexts.get(1).fireInstanceId());
assertEquals("persistent-invocation", contexts.get(1).invocationId());
} finally {
restarted.close();
}
}
private static void awaitRetryTrigger(JdbcDataSource dataSource) throws Exception {
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
do {
try (Connection connection = dataSource.getConnection();
Statement statement = connection.createStatement();
ResultSet result = statement.executeQuery(
"SELECT COUNT(*) FROM QRTZ_TRIGGERS WHERE TRIGGER_GROUP LIKE 'ea.scheduler.retry.%'")) {
if (result.next() && result.getInt(1) > 0) return;
}
Thread.sleep(25L);
} while (System.nanoTime() < deadline);
fail("persistent retry trigger was not stored");
}
private static JdbcDataSource dataSource() { private static JdbcDataSource dataSource() {
JdbcDataSource dataSource = new JdbcDataSource(); JdbcDataSource dataSource = new JdbcDataSource();
dataSource.setURL("jdbc:h2:mem:scheduler-" + UUID.randomUUID() + ";DB_CLOSE_DELAY=-1"); dataSource.setURL("jdbc:h2:mem:scheduler-" + UUID.randomUUID() + ";DB_CLOSE_DELAY=-1");
@@ -259,6 +339,8 @@ public class QuartzJdbcIntegrationTest {
false, false,
2, 2,
Thread.NORM_PRIORITY, Thread.NORM_PRIORITY,
1,
0L,
15_000L, 15_000L,
1_000L, 1_000L,
true, true,

View File

@@ -0,0 +1,105 @@
package com.easyagents.scheduler.quartz;
import com.easyagents.scheduler.ConcurrencyPolicy;
import com.easyagents.scheduler.MisfirePolicy;
import com.easyagents.scheduler.OnceSchedulePlan;
import com.easyagents.scheduler.ScheduleDefinition;
import com.easyagents.scheduler.ScheduleHandler;
import com.easyagents.scheduler.ScheduleId;
import com.easyagents.scheduler.ScheduleRefireException;
import org.junit.Test;
import org.quartz.JobDetail;
import org.quartz.JobExecutionContext;
import org.quartz.JobExecutionException;
import org.quartz.Scheduler;
import org.quartz.SchedulerException;
import org.quartz.Trigger;
import java.lang.reflect.Proxy;
import java.time.Duration;
import java.time.Instant;
import java.util.Date;
import java.util.List;
import java.util.Map;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
/** {@link QuartzRuntime} 失败控制语义测试。 */
public class QuartzRuntimeTest {
/** JobStore 无法保存延迟触发时不得把当前 fire 当作正常完成。 */
@Test
public void shouldRefireInPlaceWhenPersistentRetryCannotBeStored() throws Exception {
ScheduleDefinition definition = new ScheduleDefinition(
new ScheduleId("test", "retry-store-failure"),
"handler",
new OnceSchedulePlan(Instant.parse("2099-01-01T00:00:00Z")),
MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW,
true,
Map.of(),
"retry store failure"
);
JobDetail job = QuartzScheduleMapper.toJobDetail(definition);
Scheduler scheduler = proxy(Scheduler.class, (proxy, method, arguments) -> {
if ("scheduleJob".equals(method.getName())
&& arguments != null && arguments.length == 1
&& arguments[0] instanceof Trigger) {
throw new SchedulerException("job store unavailable");
}
if ("isShutdown".equals(method.getName())) return false;
return defaultValue(method.getReturnType());
});
Date fireTime = new Date();
JobExecutionContext context = proxy(JobExecutionContext.class,
(proxy, method, arguments) -> switch (method.getName()) {
case "getMergedJobDataMap" -> job.getJobDataMap();
case "getJobDetail" -> job;
case "getScheduler" -> scheduler;
case "getFireTime", "getScheduledFireTime" -> fireTime;
case "getFireInstanceId" -> "fire-1";
case "isRecovering" -> false;
case "getRefireCount" -> 0;
default -> defaultValue(method.getReturnType());
});
QuartzRuntime runtime = new QuartzRuntime(List.of(new ScheduleHandler() {
@Override
public String code() {
return "handler";
}
@Override
public void execute(com.easyagents.scheduler.ScheduleFireContext context) {
throw new ScheduleRefireException("registration unavailable",
new IllegalStateException("database unavailable"), 0,
Duration.ofMillis(1));
}
}), List.of());
try {
runtime.execute(context);
fail("expected refire request");
} catch (JobExecutionException exception) {
assertTrue(exception.refireImmediately());
}
}
@SuppressWarnings("unchecked")
private static <T> T proxy(Class<T> type, java.lang.reflect.InvocationHandler handler) {
return (T) Proxy.newProxyInstance(type.getClassLoader(), new Class<?>[]{type}, handler);
}
private static Object defaultValue(Class<?> type) {
if (!type.isPrimitive()) return null;
if (type == boolean.class) return false;
if (type == byte.class) return (byte) 0;
if (type == short.class) return (short) 0;
if (type == int.class) return 0;
if (type == long.class) return 0L;
if (type == float.class) return 0F;
if (type == double.class) return 0D;
if (type == char.class) return '\0';
return null;
}
}

View File

@@ -10,20 +10,25 @@ import com.easyagents.scheduler.ScheduleException;
import com.easyagents.scheduler.ScheduleFireContext; import com.easyagents.scheduler.ScheduleFireContext;
import com.easyagents.scheduler.ScheduleHandler; import com.easyagents.scheduler.ScheduleHandler;
import com.easyagents.scheduler.ScheduleId; import com.easyagents.scheduler.ScheduleId;
import com.easyagents.scheduler.ScheduleRefireException;
import com.easyagents.scheduler.ScheduleStatus; import com.easyagents.scheduler.ScheduleStatus;
import org.junit.After; import org.junit.After;
import org.junit.Test; import org.junit.Test;
import org.quartz.JobDetail; import org.quartz.JobDetail;
import org.quartz.Scheduler; import org.quartz.Scheduler;
import org.quartz.impl.StdSchedulerFactory; import org.quartz.impl.StdSchedulerFactory;
import org.quartz.impl.matchers.GroupMatcher;
import java.nio.charset.StandardCharsets; import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.Instant; import java.time.Instant;
import java.time.ZoneId; import java.time.ZoneId;
import java.util.Base64; import java.util.Base64;
import java.util.List;
import java.util.Map; import java.util.Map;
import java.util.Properties; import java.util.Properties;
import java.util.UUID; import java.util.UUID;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch; import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicInteger;
@@ -112,6 +117,108 @@ public class QuartzScheduleServiceTest {
assertEquals("value", captured.get().parameters().get("stable")); assertEquals("value", captured.get().parameters().get("stable"));
} }
/** 验证可恢复失败会持久化延迟重试,并保留原始 fire 上下文。 */
@Test
public void shouldRefireRetryableHandlerFailure() throws Exception {
AtomicInteger attempts = new AtomicInteger();
CountDownLatch completed = new CountDownLatch(1);
CopyOnWriteArrayList<ScheduleFireContext> contexts = new CopyOnWriteArrayList<>();
service = newRamService(context -> {
contexts.add(context);
if (attempts.incrementAndGet() == 1) {
throw new ScheduleRefireException("temporary registration failure",
new IllegalStateException("database unavailable"), 2,
Duration.ofMillis(100));
}
completed.countDown();
});
ScheduleDefinition definition = onceDefinition(
"retryable-immediate",
MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW,
Instant.parse("2099-01-01T00:00:00Z")
);
service.create(definition);
service.triggerNow(definition.id(), "retryable-invocation", Map.of());
assertTrue("retryable handler was not refired", completed.await(5, TimeUnit.SECONDS));
assertEquals(2, attempts.get());
assertEquals(contexts.get(0).scheduledFireTime(), contexts.get(1).scheduledFireTime());
assertEquals(contexts.get(0).actualFireTime(), contexts.get(1).actualFireTime());
assertEquals(contexts.get(0).fireInstanceId(), contexts.get(1).fireInstanceId());
assertEquals("retryable-invocation", contexts.get(1).invocationId());
}
/** 持久化延迟重试不得占用当前 Quartz worker。 */
@Test
public void shouldReleaseWorkerWhilePersistentRetryIsDelayed() throws Exception {
CountDownLatch retryScheduled = new CountDownLatch(1);
CountDownLatch healthyCompleted = new CountDownLatch(1);
CountDownLatch retryCompleted = new CountDownLatch(1);
AtomicInteger retryAttempts = new AtomicInteger();
service = newRamService(context -> {
if (context.scheduleId().name().equals("delayed-retry")) {
if (retryAttempts.incrementAndGet() == 1) {
retryScheduled.countDown();
throw new ScheduleRefireException("temporary registration failure",
new IllegalStateException("database unavailable"), 1,
Duration.ofSeconds(1));
}
retryCompleted.countDown();
} else {
healthyCompleted.countDown();
}
}, 30_000L, 1);
ScheduleDefinition retry = onceDefinition(
"delayed-retry", MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW, Instant.parse("2099-01-01T00:00:00Z"));
ScheduleDefinition healthy = onceDefinition(
"healthy-during-retry", MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW, Instant.parse("2099-01-01T00:00:00Z"));
service.create(retry);
service.create(healthy);
service.triggerNow(retry.id(), "retry-1", Map.of());
assertTrue(retryScheduled.await(5, TimeUnit.SECONDS));
service.triggerNow(healthy.id(), "healthy-1", Map.of());
assertTrue("single Quartz worker remained occupied by delayed retry",
healthyCompleted.await(750, TimeUnit.MILLISECONDS));
assertTrue(retryCompleted.await(5, TimeUnit.SECONDS));
}
/** 已持久化 fire 的重试不得因 replace 而切换到新 Handler。 */
@Test
public void shouldKeepOriginalHandlerWhenScheduleIsReplacedDuringRetry() throws Exception {
AtomicInteger oldAttempts = new AtomicInteger();
AtomicInteger newAttempts = new AtomicInteger();
CountDownLatch oldRetryCompleted = new CountDownLatch(1);
service = newRamServiceWithHandlers(List.of(
handler("old-handler", context -> {
if (oldAttempts.incrementAndGet() == 1) {
throw new ScheduleRefireException("temporary registration failure",
new IllegalStateException("database unavailable"), 1,
Duration.ofMillis(400));
}
oldRetryCompleted.countDown();
}),
handler("new-handler", context -> newAttempts.incrementAndGet())
), 30_000L, 2);
ScheduleId id = new ScheduleId("test", "replace-during-retry");
ScheduleDefinition original = definition(id, "old-handler");
service.create(original);
service.triggerNow(id, "replace-retry-1", Map.of());
awaitPersistentRetryTrigger();
service.replace(definition(id, "new-handler"));
assertTrue("old handler retry did not complete",
oldRetryCompleted.await(5, TimeUnit.SECONDS));
assertEquals(2, oldAttempts.get());
assertEquals(0, newAttempts.get());
}
/** /**
* 验证立即触发会在返回回执前校验基础参数与覆盖参数的合并结果。 * 验证立即触发会在返回回执前校验基础参数与覆盖参数的合并结果。
*/ */
@@ -385,6 +492,23 @@ public class QuartzScheduleServiceTest {
private QuartzScheduleService newRamService( private QuartzScheduleService newRamService(
Consumer<ScheduleFireContext> handler, Consumer<ScheduleFireContext> handler,
long shutdownWaitTimeoutMillis long shutdownWaitTimeoutMillis
) throws Exception {
return newRamService(handler, shutdownWaitTimeoutMillis, 2);
}
private QuartzScheduleService newRamService(
Consumer<ScheduleFireContext> handler,
long shutdownWaitTimeoutMillis,
int threadCount
) throws Exception {
return newRamServiceWithHandlers(List.of(handler("handler", handler)),
shutdownWaitTimeoutMillis, threadCount);
}
private QuartzScheduleService newRamServiceWithHandlers(
List<ScheduleHandler> handlers,
long shutdownWaitTimeoutMillis,
int threadCount
) throws Exception { ) throws Exception {
Properties properties = new Properties(); Properties properties = new Properties();
properties.setProperty( properties.setProperty(
@@ -394,7 +518,7 @@ public class QuartzScheduleServiceTest {
properties.setProperty("org.quartz.scheduler.instanceId", "NON_CLUSTERED"); properties.setProperty("org.quartz.scheduler.instanceId", "NON_CLUSTERED");
properties.setProperty("org.quartz.scheduler.interruptJobsOnShutdown", "true"); properties.setProperty("org.quartz.scheduler.interruptJobsOnShutdown", "true");
properties.setProperty("org.quartz.threadPool.class", "org.quartz.simpl.SimpleThreadPool"); properties.setProperty("org.quartz.threadPool.class", "org.quartz.simpl.SimpleThreadPool");
properties.setProperty("org.quartz.threadPool.threadCount", "2"); properties.setProperty("org.quartz.threadPool.threadCount", Integer.toString(threadCount));
properties.setProperty("org.quartz.jobStore.class", "org.quartz.simpl.RAMJobStore"); properties.setProperty("org.quartz.jobStore.class", "org.quartz.simpl.RAMJobStore");
properties.setProperty("org.quartz.jobStore.misfireThreshold", "100"); properties.setProperty("org.quartz.jobStore.misfireThreshold", "100");
Scheduler scheduler = new StdSchedulerFactory(properties).getScheduler(); Scheduler scheduler = new StdSchedulerFactory(properties).getScheduler();
@@ -402,17 +526,7 @@ public class QuartzScheduleServiceTest {
scheduler, scheduler,
true, true,
shutdownWaitTimeoutMillis, shutdownWaitTimeoutMillis,
java.util.List.of(new ScheduleHandler() { handlers,
@Override
public String code() {
return "handler";
}
@Override
public void execute(ScheduleFireContext context) {
handler.accept(context);
}
}),
java.util.List.of() java.util.List.of()
); );
result.start(); result.start();
@@ -420,6 +534,50 @@ public class QuartzScheduleServiceTest {
return result; return result;
} }
private void awaitPersistentRetryTrigger() throws Exception {
long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
while (System.nanoTime() < deadline) {
boolean found = service.quartzScheduler()
.getTriggerKeys(GroupMatcher.anyTriggerGroup())
.stream()
.anyMatch(key -> key.getGroup().startsWith(
QuartzScheduleMapper.GROUP_PREFIX + "retry."));
if (found) return;
Thread.sleep(10L);
}
fail("persistent retry trigger was not created");
}
private static ScheduleHandler handler(
String code,
Consumer<ScheduleFireContext> consumer
) {
return new ScheduleHandler() {
@Override
public String code() {
return code;
}
@Override
public void execute(ScheduleFireContext context) {
consumer.accept(context);
}
};
}
private static ScheduleDefinition definition(ScheduleId id, String handlerCode) {
return new ScheduleDefinition(
id,
handlerCode,
new OnceSchedulePlan(Instant.parse("2099-01-01T00:00:00Z")),
MisfirePolicy.FIRE_ONCE_NOW,
ConcurrencyPolicy.DISALLOW,
true,
Map.of(),
id.name()
);
}
private static ScheduleDefinition cronDefinition(String handlerCode, String expression) { private static ScheduleDefinition cronDefinition(String handlerCode, String expression) {
return new ScheduleDefinition( return new ScheduleDefinition(
new ScheduleId("test", "lifecycle"), new ScheduleId("test", "lifecycle"),

View File

@@ -0,0 +1,40 @@
package com.easyagents.scheduler.quartz;
import org.junit.Test;
import java.util.Properties;
import static org.junit.Assert.assertEquals;
/** {@link QuartzSchedulerFactory} 原生属性映射测试。 */
public class QuartzSchedulerFactoryTest {
@Test
public void batchAcquisitionMustRunWithinJobStoreLock() {
QuartzSchedulerConfig config = new QuartzSchedulerConfig(
"batch-scheduler",
"NON_CLUSTERED",
"QRTZ_",
QuartzSchedulerConfig.STANDARD_JDBC_DELEGATE,
false,
8,
Thread.NORM_PRIORITY,
8,
1_000L,
15_000L,
60_000L,
true,
30_000L,
false
);
Properties properties = QuartzSchedulerFactory.properties(config, "testDs");
assertEquals("8", properties.getProperty(
"org.quartz.scheduler.batchTriggerAcquisitionMaxCount"));
assertEquals("1000", properties.getProperty(
"org.quartz.scheduler.batchTriggerAcquisitionFireAheadTimeWindow"));
assertEquals("true", properties.getProperty(
"org.quartz.jobStore.acquireTriggersWithinLock"));
}
}

View File

@@ -53,6 +53,11 @@
<artifactId>h2</artifactId> <artifactId>h2</artifactId>
<scope>test</scope> <scope>test</scope>
</dependency> </dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-jdbc</artifactId>
<scope>test</scope>
</dependency>
<dependency> <dependency>
<groupId>junit</groupId> <groupId>junit</groupId>
<artifactId>junit</artifactId> <artifactId>junit</artifactId>

View File

@@ -18,6 +18,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
import org.springframework.boot.context.properties.EnableConfigurationProperties; import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.boot.sql.init.dependency.DependsOnDatabaseInitialization;
import org.springframework.context.annotation.Bean; import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration; import org.springframework.context.annotation.Configuration;
@@ -61,6 +62,7 @@ public class EasyAgentsSchedulerAutoConfiguration {
*/ */
@Bean(name = SCHEDULER_BEAN_NAME, initMethod = "start", destroyMethod = "close") @Bean(name = SCHEDULER_BEAN_NAME, initMethod = "start", destroyMethod = "close")
@ConditionalOnMissingBean(ScheduleService.class) @ConditionalOnMissingBean(ScheduleService.class)
@DependsOnDatabaseInitialization
public QuartzScheduleService easyAgentsQuartzScheduleService( public QuartzScheduleService easyAgentsQuartzScheduleService(
EasyAgentsSchedulerProperties properties, EasyAgentsSchedulerProperties properties,
ListableBeanFactory beanFactory, ListableBeanFactory beanFactory,
@@ -83,6 +85,8 @@ public class EasyAgentsSchedulerAutoConfiguration {
quartz.isClustered(), quartz.isClustered(),
quartz.getThreadCount(), quartz.getThreadCount(),
quartz.getThreadPriority(), quartz.getThreadPriority(),
quartz.getBatchTriggerAcquisitionMaxCount(),
quartz.getBatchTriggerAcquisitionFireAheadTimeWindowMillis(),
quartz.getClusterCheckinIntervalMillis(), quartz.getClusterCheckinIntervalMillis(),
quartz.getMisfireThresholdMillis(), quartz.getMisfireThresholdMillis(),
quartz.isWaitForJobsToCompleteOnShutdown(), quartz.isWaitForJobsToCompleteOnShutdown(),

View File

@@ -104,6 +104,12 @@ public class EasyAgentsSchedulerProperties {
/** Quartz Worker 线程优先级。 */ /** Quartz Worker 线程优先级。 */
private int threadPriority = Thread.NORM_PRIORITY; private int threadPriority = Thread.NORM_PRIORITY;
/** 单次批量获取 Trigger 的最大数量。 */
private int batchTriggerAcquisitionMaxCount = 1;
/** 可提前纳入批量的时间窗口,单位毫秒。 */
private long batchTriggerAcquisitionFireAheadTimeWindowMillis;
/** 集群心跳间隔,单位毫秒。 */ /** 集群心跳间隔,单位毫秒。 */
private long clusterCheckinIntervalMillis = 15_000L; private long clusterCheckinIntervalMillis = 15_000L;
@@ -251,6 +257,45 @@ public class EasyAgentsSchedulerProperties {
this.threadPriority = threadPriority; this.threadPriority = threadPriority;
} }
/**
* 返回单次批量获取 Trigger 的最大数量。
*
* @return 批量上限
*/
public int getBatchTriggerAcquisitionMaxCount() {
return batchTriggerAcquisitionMaxCount;
}
/**
* 设置单次批量获取 Trigger 的最大数量。
*
* @param batchTriggerAcquisitionMaxCount 批量上限
*/
public void setBatchTriggerAcquisitionMaxCount(
int batchTriggerAcquisitionMaxCount) {
this.batchTriggerAcquisitionMaxCount = batchTriggerAcquisitionMaxCount;
}
/**
* 返回可提前纳入批量的时间窗口。
*
* @return 毫秒窗口
*/
public long getBatchTriggerAcquisitionFireAheadTimeWindowMillis() {
return batchTriggerAcquisitionFireAheadTimeWindowMillis;
}
/**
* 设置可提前纳入批量的时间窗口。
*
* @param batchTriggerAcquisitionFireAheadTimeWindowMillis 毫秒窗口
*/
public void setBatchTriggerAcquisitionFireAheadTimeWindowMillis(
long batchTriggerAcquisitionFireAheadTimeWindowMillis) {
this.batchTriggerAcquisitionFireAheadTimeWindowMillis =
batchTriggerAcquisitionFireAheadTimeWindowMillis;
}
/** /**
* 返回集群心跳间隔。 * 返回集群心跳间隔。
* *

View File

@@ -114,6 +114,47 @@ public class EasyAgentsSchedulerAutoConfigurationTest {
} }
} }
/**
* 验证调度器等待 Spring Boot 数据库脚本初始化完成后再启动。
*/
@Test
public void shouldWaitForDatabaseInitializationBeforeStartingScheduler() throws Exception {
JdbcDataSource dataSource = dataSource();
Map<String, Object> properties = enabledProperties();
properties.put("easy-agents.scheduler.data-source-bean-name", "schedulerDataSource");
properties.put("spring.sql.init.mode", "always");
properties.put(
"spring.sql.init.schema-locations",
"classpath:quartz-schema/h2-2.5.2.sql"
);
SpringApplication application = new SpringApplication(AutoDiscoveryApplication.class);
application.setWebApplicationType(WebApplicationType.NONE);
application.setDefaultProperties(properties);
application.addInitializers(applicationContext -> {
GenericApplicationContext genericContext =
(GenericApplicationContext) applicationContext;
genericContext.registerBean(
"schedulerDataSource",
DataSource.class,
() -> dataSource
);
});
try (ConfigurableApplicationContext context = application.run()) {
assertNotNull(context.getBean(ScheduleService.class));
try (
Connection connection = dataSource.getConnection();
Statement statement = connection.createStatement();
ResultSet resultSet = statement.executeQuery(
"SELECT COUNT(*) FROM QRTZ_SCHEDULER_STATE"
)
) {
assertTrue(resultSet.next());
}
}
}
/** /**
* 验证多个 DataSource 未明确选择时启动失败并提供可操作信息。 * 验证多个 DataSource 未明确选择时启动失败并提供可操作信息。
*/ */
@@ -162,6 +203,13 @@ public class EasyAgentsSchedulerAutoConfigurationTest {
properties.put("easy-agents.scheduler.quartz.instance-id", "NON_CLUSTERED"); properties.put("easy-agents.scheduler.quartz.instance-id", "NON_CLUSTERED");
properties.put("easy-agents.scheduler.quartz.clustered", "false"); properties.put("easy-agents.scheduler.quartz.clustered", "false");
properties.put("easy-agents.scheduler.quartz.thread-count", "2"); properties.put("easy-agents.scheduler.quartz.thread-count", "2");
properties.put(
"easy-agents.scheduler.quartz.batch-trigger-acquisition-max-count", "2"
);
properties.put(
"easy-agents.scheduler.quartz.batch-trigger-acquisition-fire-ahead-time-window-millis",
"500"
);
properties.put("easy-agents.scheduler.quartz.misfire-threshold-millis", "1000"); properties.put("easy-agents.scheduler.quartz.misfire-threshold-millis", "1000");
return properties; return properties;
} }

View File

@@ -199,10 +199,13 @@ public class LuceneSearcher implements DocumentSearcher, AutoCloseable {
@Override @Override
public List<Document> searchDocuments(KeywordSearchRequest request) { public List<Document> searchDocuments(KeywordSearchRequest request) {
List<Document> results = new ArrayList<>(); List<Document> results = new ArrayList<>();
if (request == null || request.getKeyword() == null || request.getKeyword().trim().isEmpty()) {
return results;
}
try (IndexReader reader = DirectoryReader.open(directory)) { try (IndexReader reader = DirectoryReader.open(directory)) {
IndexSearcher searcher = new IndexSearcher(reader); IndexSearcher searcher = new IndexSearcher(reader);
Query query = buildQuery(request); Query query = buildQuery(request);
TopDocs topDocs = searcher.search(query, request == null ? 10 : request.getCount()); TopDocs topDocs = searcher.search(query, request.getCount());
for (ScoreDoc scoreDoc : topDocs.scoreDocs) { for (ScoreDoc scoreDoc : topDocs.scoreDocs) {
org.apache.lucene.document.Document doc = searcher.doc(scoreDoc.doc); org.apache.lucene.document.Document doc = searcher.doc(scoreDoc.doc);
Document resultDoc = new Document(); Document resultDoc = new Document();
@@ -224,29 +227,24 @@ public class LuceneSearcher implements DocumentSearcher, AutoCloseable {
return results; return results;
} }
Query buildQuery(KeywordSearchRequest request) { Query buildQuery(KeywordSearchRequest request) throws ParseException {
try { String escapedKeyword = QueryParser.escape(request.getKeyword());
String keyword = request == null ? null : request.getKeyword();
QueryParser titleQueryParser = new QueryParser("title", analyzer); QueryParser titleQueryParser = new QueryParser("title", analyzer);
Query titleQuery = titleQueryParser.parse(keyword); Query titleQuery = titleQueryParser.parse(escapedKeyword);
BooleanClause titleBooleanClause = new BooleanClause(titleQuery, BooleanClause.Occur.SHOULD); BooleanClause titleBooleanClause = new BooleanClause(titleQuery, BooleanClause.Occur.SHOULD);
QueryParser contentQueryParser = new QueryParser("content", analyzer); QueryParser contentQueryParser = new QueryParser("content", analyzer);
Query contentQuery = contentQueryParser.parse(keyword); Query contentQuery = contentQueryParser.parse(escapedKeyword);
BooleanClause contentBooleanClause = new BooleanClause(contentQuery, BooleanClause.Occur.SHOULD); BooleanClause contentBooleanClause = new BooleanClause(contentQuery, BooleanClause.Occur.SHOULD);
BooleanQuery.Builder builder = new BooleanQuery.Builder(); BooleanQuery.Builder builder = new BooleanQuery.Builder();
builder.add(titleBooleanClause) builder.add(titleBooleanClause)
.add(contentBooleanClause); .add(contentBooleanClause);
if (request != null && request.getKnowledgeId() != null && !request.getKnowledgeId().trim().isEmpty()) { if (request.getKnowledgeId() != null && !request.getKnowledgeId().trim().isEmpty()) {
builder.add(new TermQuery(new Term(KeywordSearchMetadataKeys.KNOWLEDGE_ID, request.getKnowledgeId().trim())), BooleanClause.Occur.MUST); builder.add(new TermQuery(new Term(KeywordSearchMetadataKeys.KNOWLEDGE_ID, request.getKnowledgeId().trim())), BooleanClause.Occur.MUST);
} }
return builder.build(); return builder.build();
} catch (ParseException e) {
LOG.error(e.toString(), e);
}
return null;
} }
private static Analyzer createAnalyzer() { private static Analyzer createAnalyzer() {

View File

@@ -57,6 +57,45 @@ public class LuceneSearcherTest {
} }
} }
/**
* 验证用户输入中的 Lucene 特殊字符按普通文本检索。
*
* @throws Exception 临时目录或 Lucene 资源操作失败时抛出
*/
@Test
public void shouldTreatLuceneSpecialCharactersAsPlainText() throws Exception {
Path tempDir = Files.createTempDirectory("lucene-searcher-special-character-test");
LuceneConfig config = new LuceneConfig();
config.setIndexDirPath(tempDir.toString());
try (LuceneSearcher searcher = new LuceneSearcher(config)) {
Document document = new Document();
document.setId("special-character");
document.setContent("A/C.\\nPLEASE");
Assert.assertTrue(searcher.addDocument(document));
List<Document> results = searcher.searchDocuments("A/C.\\nPLEASE", 10);
Assert.assertEquals(1, results.size());
Assert.assertEquals("special-character", String.valueOf(results.get(0).getId()));
}
}
/**
* 验证空查询直接返回空结果。
*
* @throws Exception 临时目录或 Lucene 资源操作失败时抛出
*/
@Test
public void shouldReturnEmptyForMissingKeyword() throws Exception {
Path tempDir = Files.createTempDirectory("lucene-searcher-empty-keyword-test");
LuceneConfig config = new LuceneConfig();
config.setIndexDirPath(tempDir.toString());
try (LuceneSearcher searcher = new LuceneSearcher(config)) {
Assert.assertTrue(searcher.searchDocuments((KeywordSearchRequest) null).isEmpty());
Assert.assertTrue(searcher.searchDocuments(" ", 10).isEmpty());
}
}
/** /**
* 验证多个导入线程可共享同一个 IndexWriter 完成批量写入。 * 验证多个导入线程可共享同一个 IndexWriter 完成批量写入。
* *

View File

@@ -25,7 +25,6 @@
<dependency> <dependency>
<groupId>io.milvus</groupId> <groupId>io.milvus</groupId>
<artifactId>milvus-sdk-java</artifactId> <artifactId>milvus-sdk-java</artifactId>
<version>2.4.1</version>
</dependency> </dependency>
<dependency> <dependency>
<groupId>junit</groupId> <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; package com.easyagents.store.milvus;
import com.alibaba.fastjson.JSON; import com.google.gson.Gson;
import com.alibaba.fastjson.JSONObject; import com.google.gson.JsonObject;
import com.easyagents.core.document.Document; import com.easyagents.core.document.Document;
import com.easyagents.core.store.DocumentStore; import com.easyagents.core.store.DocumentStore;
import com.easyagents.core.store.SearchWrapper; import com.easyagents.core.store.SearchWrapper;
import com.easyagents.core.store.StoreOptions; import com.easyagents.core.store.StoreOptions;
import com.easyagents.core.store.StoreResult; import com.easyagents.core.store.StoreResult;
import com.easyagents.core.store.StoreTimeoutException;
import com.easyagents.core.util.CollectionUtil; import com.easyagents.core.util.CollectionUtil;
import com.easyagents.core.util.Maps; import com.easyagents.core.util.Maps;
import com.easyagents.core.util.StringUtil; 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.client.MilvusClientV2;
import io.milvus.v2.common.ConsistencyLevel; import io.milvus.v2.common.ConsistencyLevel;
import io.milvus.v2.common.DataType; import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam; 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.CreateCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq; import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.request.HasCollectionReq; import io.milvus.v2.service.collection.request.HasCollectionReq;
import io.milvus.v2.service.collection.request.LoadCollectionReq; import io.milvus.v2.service.collection.request.LoadCollectionReq;
import io.milvus.v2.service.vector.request.*; 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.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp; import io.milvus.v2.service.vector.response.SearchResp;
import org.slf4j.Logger; import org.slf4j.Logger;
import org.slf4j.LoggerFactory; import org.slf4j.LoggerFactory;
import java.net.URI; import java.time.Duration;
import java.util.*; 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. * 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 Logger LOG = LoggerFactory.getLogger(MilvusVectorStore.class);
private static final long LOAD_TIMEOUT_MS = 30_000L; private static final long LOAD_TIMEOUT_MS = 30_000L;
private static final long LOAD_POLL_INTERVAL_MS = 200L; 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_ID = "id";
private static final String FIELD_CONTENT = "content"; private static final String FIELD_CONTENT = "content";
private static final String FIELD_METADATA = "metadata"; private static final String FIELD_METADATA = "metadata";
private static final String FIELD_VECTOR = "vector"; 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 MilvusVectorStoreConfig config;
private final String defaultCollectionName; private final String defaultCollectionName;
private final Set<String> initializedCollections = Collections.synchronizedSet(new HashSet<String>()); private final boolean ownsClientManager;
private final Set<String> loadedCollections = Collections.synchronizedSet(new HashSet<String>()); private volatile MilvusClientV2 compatibilityClient;
private volatile boolean closed;
public MilvusVectorStore(MilvusVectorStoreConfig config) { 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.config = config;
this.defaultCollectionName = config.getDefaultCollectionName(); this.defaultCollectionName = config.getDefaultCollectionName();
String uri = normalizeAndValidateUri(config.getUri()); this.clientManager = Objects.requireNonNull(clientManager, "clientManager");
String dbName = StringUtil.hasText(config.getDatabaseName()) ? config.getDatabaseName().trim() : "default"; this.ownsClientManager = ownsClientManager;
ConnectConfig.ConnectConfigBuilder<?, ?> builder = ConnectConfig.builder()
.uri(uri)
.dbName(dbName);
if (StringUtil.hasText(config.getToken())) {
builder.token(config.getToken().trim());
} }
if (StringUtil.hasText(config.getUsername()) && StringUtil.hasText(config.getPassword())) { private static MilvusClientManager createOwnedClientManager(
builder.username(config.getUsername().trim()); MilvusVectorStoreConfig config
builder.password(config.getPassword().trim()); ) {
validateConfig(config);
return new MilvusClientManager(config);
} }
ConnectConfig connectConfig = builder.build(); private static void validateConfig(MilvusVectorStoreConfig config) {
this.client = new MilvusClientV2(connectConfig); Objects.requireNonNull(config, "config");
if (config.getSearchTimeoutMillis() <= DEADLINE_SAFETY_MARGIN_MS) {
throw new IllegalArgumentException(
"Milvus searchTimeoutMillis must be greater than "
+ DEADLINE_SAFETY_MARGIN_MS
);
} }
if (config.getPoolMaxWaitMillis() <= 0L) {
private String normalizeAndValidateUri(String uri) { throw new IllegalArgumentException(
if (StringUtil.noText(uri)) { "Milvus poolMaxWaitMillis must be greater than zero"
throw new IllegalArgumentException("Milvus uri is required. Example: http://127.0.0.1:19530"); );
} }
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 @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."); 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 { try {
int dimension = getDimension(documents);
clientManager.withClient(client -> {
ensureCollectionExists(client, collectionName, dimension);
InsertReq.InsertReqBuilder<?, ?> builder = InsertReq.builder(); InsertReq.InsertReqBuilder<?, ?> builder = InsertReq.builder();
if (StringUtil.hasText(options.getPartitionName())) { if (StringUtil.hasText(options.getPartitionName())) {
builder.partitionName(options.getPartitionName()); builder.partitionName(options.getPartitionName());
} }
InsertReq insertReq = builder client.insert(builder
.collectionName(collectionName) .collectionName(collectionName)
.data(toMilvusDocuments(documents)) .data(toMilvusDocuments(documents))
.build(); .build());
client.insert(insertReq); return null;
});
return StoreResult.successWithIds(documents); return StoreResult.successWithIds(documents);
} catch (MilvusClientException e) { } catch (RuntimeException e) {
return StoreResult.fail(); 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) .collectionName(collectionName)
.ids(MilvusPrimaryKeySupport.normalize(ids)) .ids(MilvusPrimaryKeySupport.normalize(ids))
.build(); .build();
clientManager.withClient(client -> {
client.delete(deleteReq); client.delete(deleteReq);
return null;
});
return StoreResult.success(); return StoreResult.success();
} catch (Exception e) { } catch (Exception e) {
LOG.error("Milvus delete failed. collection={}, message={}", 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."); 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 { try {
UpsertReq upsertReq = UpsertReq.builder() int dimension = getDimension(documents);
clientManager.withClient(client -> {
ensureCollectionExists(client, collectionName, dimension);
client.upsert(UpsertReq.builder()
.collectionName(collectionName) .collectionName(collectionName)
.partitionName(options.getPartitionName()) .partitionName(options.getPartitionName())
.data(toMilvusDocuments(documents)) .data(toMilvusDocuments(documents))
.build(); .build());
client.upsert(upsertReq); return null;
});
return StoreResult.successWithIds(documents); return StoreResult.successWithIds(documents);
} catch (Exception e) { } 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)) { if (StringUtil.noText(collectionName)) {
throw new IllegalStateException("CollectionName is null or blank. please config the \"defaultCollectionName\" or store with designative 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) { 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); SearchReq searchReq = buildSearchReq(wrapper, options, collectionName);
try { SearchResp resp = withClientBeforeDeadline(
SearchResp resp = client.search(searchReq); deadlineNanos, client -> client.search(searchReq));
return parseSearchResults(resp, wrapper.getMinScore()); 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); QueryReq queryReq = buildQueryReq(wrapper, options, collectionName);
try { QueryResp resp = withClientBeforeDeadline(
QueryResp resp = client.query(queryReq); deadlineNanos, client -> client.query(queryReq));
return parseQueryResults(resp); 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) { private SearchReq buildSearchReq(SearchWrapper wrapper, StoreOptions options, String collectionName) {
@@ -259,7 +338,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
.outputFields(getOutputFields(wrapper)) .outputFields(getOutputFields(wrapper))
.topK(wrapper.getMaxResults()) .topK(wrapper.getMaxResults())
.annsField(FIELD_VECTOR) .annsField(FIELD_VECTOR)
.data(Collections.singletonList(toFloatList(wrapper.getVector()))) .data(Collections.singletonList(new FloatVec(wrapper.getVector())))
.searchParams(Maps.of("ef", 64)); .searchParams(Maps.of("ef", 64));
if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) { if (CollectionUtil.hasItems(options.getPartitionNamesOrEmpty())) {
@@ -305,11 +384,7 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
continue; continue;
} }
document.setId(result.getId()); document.setId(result.getId());
Float distance = result.getDistance(); document.setScore(normalizeScore(result.getScore()));
if (distance != null) {
double score = (distance + 1.0d) / 2.0d;
document.setScore(score);
}
if (minScore == null || document.getScore() == null || document.getScore() >= minScore) { if (minScore == null || document.getScore() == null || document.getScore() >= minScore) {
documents.add(document); documents.add(document);
} }
@@ -318,6 +393,10 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
return documents; return documents;
} }
static Double normalizeScore(Float rawScore) {
return rawScore == null ? null : (rawScore + 1.0d) / 2.0d;
}
private List<Document> parseQueryResults(QueryResp resp) { private List<Document> parseQueryResults(QueryResp resp) {
List<QueryResp.QueryResult> results = resp.getQueryResults(); List<QueryResp.QueryResult> results = resp.getQueryResults();
if (CollectionUtil.noItems(results)) { if (CollectionUtil.noItems(results)) {
@@ -360,22 +439,28 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
document.addMetadata(metadata); document.addMetadata(metadata);
} else if (metadataObj != null) { } else if (metadataObj != null) {
@SuppressWarnings("unchecked") @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); document.addMetadata(metadata);
} }
return document; return document;
} }
private List<JSONObject> toMilvusDocuments(List<Document> documents) { List<JsonObject> toMilvusDocuments(List<Document> documents) {
List<JSONObject> rows = new ArrayList<JSONObject>(documents.size()); List<JsonObject> rows = new ArrayList<JsonObject>(documents.size());
for (Document doc : documents) { for (Document doc : documents) {
JSONObject row = new JSONObject(); JsonObject row = new JsonObject();
row.put(FIELD_ID, String.valueOf(doc.getId())); row.addProperty(FIELD_ID, String.valueOf(doc.getId()));
row.put(FIELD_CONTENT, doc.getContent()); row.addProperty(FIELD_CONTENT, doc.getContent());
row.put(FIELD_VECTOR, toFloatList(doc.getVector())); row.add(FIELD_VECTOR, GSON.toJsonTree(toFloatList(doc.getVector())));
Map<String, Object> metadatas = doc.getMetadataMap(); 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); rows.add(row);
} }
return rows; return rows;
@@ -413,67 +498,175 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
throw new IllegalStateException("Unable to determine vector dimension for Milvus collection."); throw new IllegalStateException("Unable to determine vector dimension for Milvus collection.");
} }
private void ensureCollectionExists(String collectionName, int dimension) { private void ensureCollectionExists(
if (initializedCollections.contains(collectionName)) { MilvusClientV2 client,
String collectionName,
int dimension
) {
if (clientManager.isCollectionInitialized(collectionName)) {
return; return;
} }
synchronized (initializedCollections) { synchronized (clientManager.initializedCollectionsLock()) {
if (initializedCollections.contains(collectionName)) { if (clientManager.isCollectionInitialized(collectionName)) {
return; return;
} }
Boolean exists = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build()); Boolean exists = client.hasCollection(HasCollectionReq.builder().collectionName(collectionName).build());
if (Boolean.TRUE.equals(exists)) { if (Boolean.TRUE.equals(exists)) {
initializedCollections.add(collectionName); clientManager.markCollectionInitialized(collectionName);
return; return;
} }
if (!config.isAutoCreateCollection()) { if (!config.isAutoCreateCollection()) {
throw new IllegalStateException("Milvus collection not found and autoCreateCollection is disabled: " + collectionName); throw new IllegalStateException("Milvus collection not found and autoCreateCollection is disabled: " + collectionName);
} }
createCollection(collectionName, dimension); createCollection(client, collectionName, dimension);
initializedCollections.add(collectionName); clientManager.markCollectionInitialized(collectionName);
} }
} }
private void ensureCollectionLoaded(String collectionName) { private void ensureCollectionLoaded(
if (loadedCollections.contains(collectionName)) { String collectionName,
long deadlineNanos
) {
while (!clientManager.isCollectionLoaded(collectionName)) {
MilvusClientManager.CollectionLoadTicket ticket =
clientManager.beginCollectionLoad(collectionName);
if (!ticket.isLeader()) {
if (awaitCollectionLoad(
ticket, collectionName, deadlineNanos)) {
return; return;
} }
synchronized (loadedCollections) { continue;
if (loadedCollections.contains(collectionName)) {
return;
} }
boolean loaded = false;
try { try {
loaded = Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build())); if (clientManager.isCollectionLoaded(collectionName)) {
} catch (Exception e) { clientManager.completeCollectionLoad(ticket);
LOG.warn("Milvus getLoadState failed. collection={}, message={}", collectionName, e.getMessage()); return;
} }
withClientBeforeDeadline(deadlineNanos, client -> {
boolean loaded = Boolean.TRUE.equals(client.getLoadState(
GetLoadStateReq.builder()
.collectionName(collectionName)
.build()
));
if (!loaded) { if (!loaded) {
client.loadCollection(LoadCollectionReq.builder().collectionName(collectionName).build()); client.loadCollection(LoadCollectionReq.builder()
waitForCollectionLoaded(collectionName); .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) { private boolean isLeaderLocalAbort(Throwable failure) {
long deadline = System.currentTimeMillis() + LOAD_TIMEOUT_MS; if (clientManager.isClosed()) {
while (System.currentTimeMillis() < deadline) { 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()))) { if (Boolean.TRUE.equals(client.getLoadState(GetLoadStateReq.builder().collectionName(collectionName).build()))) {
return; return;
} }
try { 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) { } catch (InterruptedException e) {
Thread.currentThread().interrupt(); Thread.currentThread().interrupt();
throw new IllegalStateException("Interrupted while loading Milvus collection: " + collectionName, e); 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; Throwable current = e;
while (current != null) { while (current != null) {
String message = current.getMessage(); String message = current.getMessage();
@@ -485,7 +678,91 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
return false; 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>(); List<CreateCollectionReq.FieldSchema> fieldSchemaList = new ArrayList<CreateCollectionReq.FieldSchema>();
fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder() fieldSchemaList.add(CreateCollectionReq.FieldSchema.builder()
.name(FIELD_ID) .name(FIELD_ID)
@@ -531,31 +808,83 @@ public class MilvusVectorStore extends DocumentStore implements AutoCloseable {
.indexParams(indexParams) .indexParams(indexParams)
.build(); .build();
client.createCollection(createCollectionReq); client.createCollection(createCollectionReq);
ensureCollectionLoaded(collectionName); ensureCollectionLoadedForWrite(client, collectionName);
} }
public MilvusClientV2 getClient() { private void ensureCollectionLoadedForWrite(
return client; 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() { public boolean checkAvailable() {
try { try {
return client.hasCollection(HasCollectionReq.builder() return clientManager.withClient(client -> client.hasCollection(
HasCollectionReq.builder()
.collectionName("__milvus_boot_probe__") .collectionName("__milvus_boot_probe__")
.build()) != null; .build()
)) != null;
} catch (Exception e) { } catch (Exception e) {
LOG.warn("Milvus availability check failed. message={}", e.getMessage()); LOG.warn("Milvus availability check failed. message={}", e.getMessage());
return false; 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 @Override
public void close() { public void close() {
MilvusClientV2 legacyClient;
synchronized (this) {
if (closed) {
return;
}
closed = true;
legacyClient = compatibilityClient;
compatibilityClient = null;
}
if (legacyClient != null) {
try { try {
client.close(1L); legacyClient.close(1L);
} catch (InterruptedException e) { } catch (InterruptedException exception) {
Thread.currentThread().interrupt(); 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 password;
private String defaultCollectionName; private String defaultCollectionName;
private boolean autoCreateCollection = true; 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() { public String getUri() {
return uri; return uri;
@@ -87,6 +95,70 @@ public class MilvusVectorStoreConfig implements DocumentStoreConfig {
this.autoCreateCollection = autoCreateCollection; 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 @Override
public boolean checkAvailable() { public boolean checkAvailable() {
return StringUtil.hasText(this.uri); 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"); config.setPassword("Milvus");
Assert.assertTrue(config.checkAvailable()); 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

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