From edffc81991441eef863618d779ed8358c13f90c2 Mon Sep 17 00:00:00 2001 From: YunaiV Date: Thu, 22 Oct 2026 22:26:07 +0800 Subject: [PATCH] feat(ai1): 增加长期记忆并完善会话管理与生成控制 --- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java | 9 +++++---- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentRespVO.java | 3 +++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentSaveReqVO.java | 3 +++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java | 111 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java | 21 +++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java | 39 +++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java | 22 ++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionController.java | 13 ++++++++++--- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java | 19 +++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java | 20 ++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionPageReqVO.java | 4 +--- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionRespVO.java | 3 +++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/agent/Ai1AgentDO.java | 6 ++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java | 57 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/agent/Ai1AgentMapper.java | 2 +- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java | 32 ++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMapper.java | 12 +++++++++++- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java | 48 ++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java | 27 +++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java | 73 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java | 3 +++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java | 20 ++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java | 6 +++--- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java | 193 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java | 162 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamTool.java | 507 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-------------------------------- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java | 74 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java | 18 +++++++++++++++--- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java | 27 +++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java | 4 ++-- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImpl.java | 7 ++++++- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java | 90 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java | 123 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java | 49 +++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImpl.java | 48 ++++++++++++++++++++++++++++++++++++++++++++++-- yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java | 16 ++++++++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImpl.java | 10 ++++++++++ yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java | 52 ++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatToolTest.java | 38 ++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java | 294 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java | 149 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamToolTest.java | 672 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++--- yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionTaskRecoveryIntegrationTest.java | 135 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++------- yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java | 175 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactoryTest.java | 44 ++++++++++++++++++++++++++++++++------------ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImplTest.java | 39 +++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java | 224 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java | 170 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-------------------------- yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java | 42 ++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/util/Ai1UtilsTest.java | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++ yudao-module-ai1-springai/src/test/resources/sql/clean.sql | 1 + yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql | 19 +++++++++++++++++++ 52 files changed, 3885 insertions(+), 103 deletions(-) create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java create mode 100644 yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java create mode 100644 yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java create mode 100644 yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java create mode 100644 yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java create mode 100644 yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java index c38f70e..834eb19 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java @@ -1,7 +1,6 @@ package cn.iocoder.yudao.module.ai1.controller.admin.agent; import cn.hutool.core.collection.CollUtil; -import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; import cn.iocoder.yudao.framework.common.pojo.CommonResult; import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.common.util.collection.MapUtils; @@ -97,9 +96,11 @@ public class Ai1AgentController { } @GetMapping("/simple-list") - @Operation(summary = "获得已开启的 Agent 精简列表", description = "用于会话页选择 Agent，仅返回已开启的 Agent") - public CommonResult> getAgentSimpleList() { - List list = agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus()); + @Operation(summary = "获得 Agent 精简列表", description = "按状态筛选，不传时返回全部状态；会话页传 0 仅选择已开启的 Agent") + @Parameter(name = "status", description = "状态，参见 CommonStatusEnum 枚举：0 开启，1 关闭；不传时返回全部状态", example = "0") + public CommonResult> getAgentSimpleList( + @RequestParam(value = "status", required = false) Integer status) { + List list = agentService.getAgentListByStatus(status); Map modelMap = modelService.getModelMap(convertSet(list, Ai1AgentDO::getModelId)); return success(convertList(list, agent -> { Ai1AgentRespVO respVO = new Ai1AgentRespVO().setId(agent.getId()).setName(agent.getName()) diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentRespVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentRespVO.java index 16a95a6..5b728ae 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentRespVO.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentRespVO.java @@ -43,6 +43,9 @@ public class Ai1AgentRespVO { @Schema(description = "SKILL 编号集合", example = "[1]") private List skillIds; + @Schema(description = "是否开启长期记忆", requiredMode = Schema.RequiredMode.REQUIRED, example = "false") + private Boolean memoryEnabled; + @Schema(description = "状态", requiredMode = Schema.RequiredMode.REQUIRED, example = "0") private Integer status; diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentSaveReqVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentSaveReqVO.java index 8e724c9..488858e 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentSaveReqVO.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/vo/Ai1AgentSaveReqVO.java @@ -44,4 +44,7 @@ public class Ai1AgentSaveReqVO { @Schema(description = "SKILL 编号集合", example = "[1]") private List skillIds; + @Schema(description = "是否开启长期记忆", example = "false") + private Boolean memoryEnabled; + } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java new file mode 100644 index 0000000..7349de1 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java @@ -0,0 +1,111 @@ +package cn.iocoder.yudao.module.ai1.controller.admin.memory; + +import cn.hutool.core.collection.CollUtil; +import cn.iocoder.yudao.framework.common.pojo.CommonResult; +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.framework.common.util.collection.MapUtils; +import cn.iocoder.yudao.framework.common.util.object.BeanUtils; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryRespVO; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; +import cn.iocoder.yudao.module.system.api.user.AdminUserApi; +import cn.iocoder.yudao.module.system.api.user.dto.AdminUserRespDTO; +import io.swagger.v3.oas.annotations.Operation; +import io.swagger.v3.oas.annotations.Parameter; +import io.swagger.v3.oas.annotations.tags.Tag; +import jakarta.annotation.Resource; +import jakarta.validation.Valid; +import org.springframework.security.access.prepost.PreAuthorize; +import org.springframework.validation.annotation.Validated; +import org.springframework.web.bind.annotation.*; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success; +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertSet; + +@Tag(name = "管理后台 - AI1 长期记忆") +@RestController +@RequestMapping("/ai1/memory") +@Validated +public class Ai1MemoryController { + + @Resource + private Ai1MemoryService memoryService; + @Resource + private Ai1AgentService agentService; + + @Resource + private AdminUserApi adminUserApi; + + @PutMapping("/update") + @Operation(summary = "更新长期记忆") + @PreAuthorize("@ss.hasPermission('ai1:memory:update')") + public CommonResult updateMemory(@Valid @RequestBody Ai1MemoryUpdateReqVO updateReqVO) { + memoryService.updateMemory(updateReqVO); + return success(true); + } + + @DeleteMapping("/delete") + @Operation(summary = "删除长期记忆", description = "之后根据历史对话重新生成") + @Parameter(name = "id", description = "编号", required = true) + @PreAuthorize("@ss.hasPermission('ai1:memory:delete')") + public CommonResult deleteMemory(@RequestParam("id") Long id) { + memoryService.deleteMemory(id); + return success(true); + } + + @DeleteMapping("/delete-list") + @Operation(summary = "批量删除长期记忆", description = "之后根据历史对话重新生成") + @Parameter(name = "ids", description = "编号列表", required = true) + @PreAuthorize("@ss.hasPermission('ai1:memory:delete')") + public CommonResult deleteMemoryList(@RequestParam("ids") List ids) { + memoryService.deleteMemoryListByIds(ids); + return success(true); + } + + @GetMapping("/get") + @Operation(summary = "获得长期记忆") + @Parameter(name = "id", description = "编号", required = true, example = "1024") + @PreAuthorize("@ss.hasPermission('ai1:memory:query')") + public CommonResult getMemory(@RequestParam("id") Long id) { + Ai1MemoryDO memory = memoryService.getMemory(id); + return success(buildMemoryRespVO(memory)); + } + + @GetMapping("/page") + @Operation(summary = "获得长期记忆分页") + @PreAuthorize("@ss.hasPermission('ai1:memory:query')") + public CommonResult> getMemoryPage(@Valid Ai1MemoryPageReqVO pageReqVO) { + PageResult pageResult = memoryService.getMemoryPage(pageReqVO); + return success(new PageResult<>(buildMemoryRespVOList(pageResult.getList()), pageResult.getTotal())); + } + + // ==================== 拼接 VO ==================== + + private Ai1MemoryRespVO buildMemoryRespVO(Ai1MemoryDO memory) { + if (memory == null) { + return null; + } + return CollUtil.getFirst(buildMemoryRespVOList(Collections.singletonList(memory))); + } + + private List buildMemoryRespVOList(List list) { + if (CollUtil.isEmpty(list)) { + return Collections.emptyList(); + } + Map agentMap = agentService.getAgentMap(convertSet(list, Ai1MemoryDO::getAgentId)); + Map userMap = adminUserApi.getUserMap(convertSet(list, Ai1MemoryDO::getUserId)); + return BeanUtils.toBean(list, Ai1MemoryRespVO.class, respVO -> { + MapUtils.findAndThen(agentMap, respVO.getAgentId(), agent -> respVO.setAgentName(agent.getName())); + MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname())); + }); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java new file mode 100644 index 0000000..67012dd --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java @@ -0,0 +1,21 @@ +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo; + +import cn.iocoder.yudao.framework.common.pojo.PageParam; +import io.swagger.v3.oas.annotations.media.Schema; +import lombok.Data; +import lombok.EqualsAndHashCode; +import lombok.ToString; + +@Schema(description = "管理后台 - AI1 长期记忆分页 Request VO") +@Data +@EqualsAndHashCode(callSuper = true) +@ToString(callSuper = true) +public class Ai1MemoryPageReqVO extends PageParam { + + @Schema(description = "Agent 编号", example = "1") + private Long agentId; + + @Schema(description = "用户编号", example = "1") + private Long userId; + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java new file mode 100644 index 0000000..8561f40 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java @@ -0,0 +1,39 @@ +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo; + +import io.swagger.v3.oas.annotations.media.Schema; +import lombok.Data; + +import java.time.LocalDateTime; + +@Schema(description = "管理后台 - AI1 长期记忆 Response VO") +@Data +public class Ai1MemoryRespVO { + + @Schema(description = "编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024") + private Long id; + + @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") + private Long agentId; + + @Schema(description = "Agent 名称", example = "客服助手") + private String agentName; + + @Schema(description = "用户编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") + private Long userId; + + @Schema(description = "用户昵称", example = "芋道源码") + private String userNickname; + + @Schema(description = "记忆正文", requiredMode = Schema.RequiredMode.REQUIRED, example = "## 用户画像\n- 后端工程师") + private String content; + + @Schema(description = "已归档的最后一条消息编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "2048") + private Long lastMessageId; + + @Schema(description = "创建时间", requiredMode = Schema.RequiredMode.REQUIRED) + private LocalDateTime createTime; + + @Schema(description = "更新时间", requiredMode = Schema.RequiredMode.REQUIRED) + private LocalDateTime updateTime; + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java new file mode 100644 index 0000000..0084443 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java @@ -0,0 +1,22 @@ +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo; + +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.constraints.NotBlank; +import jakarta.validation.constraints.NotNull; +import jakarta.validation.constraints.Size; +import lombok.Data; + +@Schema(description = "管理后台 - AI1 长期记忆修改 Request VO") +@Data +public class Ai1MemoryUpdateReqVO { + + @Schema(description = "编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024") + @NotNull(message = "编号不能为空") + private Long id; + + @Schema(description = "记忆正文", requiredMode = Schema.RequiredMode.REQUIRED, example = "## 用户画像\n- 后端工程师") + @NotBlank(message = "记忆正文不能为空") + @Size(max = 65535, message = "记忆正文长度不能超过 65535 个字符") + private String content; + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionController.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionController.java index 1de68df..4664a19 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionController.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionController.java @@ -8,7 +8,9 @@ import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1Sessio import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionPageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionRespVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionUpdateMyReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; import cn.iocoder.yudao.module.system.api.user.AdminUserApi; import cn.iocoder.yudao.module.system.api.user.dto.AdminUserRespDTO; @@ -36,6 +38,8 @@ public class Ai1SessionController { @Resource private Ai1SessionService sessionService; + @Resource + private Ai1AgentService agentService; @Resource private AdminUserApi adminUserApi; @@ -78,13 +82,16 @@ public class Ai1SessionController { // ==================== 管理 ==================== @GetMapping("/page") - @Operation(summary = "获得会话分页", description = "管理端查看 Agent 下全部用户的会话") + @Operation(summary = "获得会话分页", description = "管理端查看全部用户的会话，可按 Agent 筛选") @PreAuthorize("@ss.hasPermission('ai1:session:query')") public CommonResult> getSessionPage(@Valid Ai1SessionPageReqVO pageReqVO) { PageResult pageResult = sessionService.getSessionPage(pageReqVO); + Map agentMap = agentService.getAgentMap(convertSet(pageResult.getList(), Ai1SessionDO::getAgentId)); Map userMap = adminUserApi.getUserMap(convertSet(pageResult.getList(), Ai1SessionDO::getUserId)); - return success(BeanUtils.toBean(pageResult, Ai1SessionRespVO.class, respVO -> - MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname())))); + return success(BeanUtils.toBean(pageResult, Ai1SessionRespVO.class, respVO -> { + MapUtils.findAndThen(agentMap, respVO.getAgentId(), agent -> respVO.setAgentName(agent.getName())); + MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname())); + })); } } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java index 111727b..2526abd 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java @@ -1,7 +1,9 @@ package cn.iocoder.yudao.module.ai1.controller.admin.session; import cn.iocoder.yudao.framework.common.pojo.CommonResult; +import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.common.util.object.BeanUtils; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageRespVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; @@ -63,8 +65,25 @@ public class Ai1SessionMessageController { return sessionMessageService.resumeSessionMessageStream(getLoginUserId(), messageId, lastEventId); } + @PostMapping("/stop-my") + @Operation(summary = "停止【我的】会话生成", description = "终止会话中生成中的助手消息，保留已生成的内容") + @Parameter(name = "sessionId", description = "会话编号", required = true, example = "1024") + @PreAuthorize("@ss.hasPermission('ai1:session:my')") + public CommonResult stopSessionMessageMy(@RequestParam("sessionId") Long sessionId) { + sessionMessageService.stopSessionMessageMy(getLoginUserId(), sessionId); + return success(true); + } + // ==================== 管理 ==================== + @GetMapping("/page") + @Operation(summary = "获得会话消息分页", description = "按消息编号倒序") + @PreAuthorize("@ss.hasPermission('ai1:session:query')") + public CommonResult> getSessionMessagePage(@Valid Ai1SessionMessagePageReqVO pageReqVO) { + return success(BeanUtils.toBean(sessionMessageService.getSessionMessagePage(pageReqVO), + Ai1SessionMessageRespVO.class)); + } + @GetMapping("/list") @Operation(summary = "获得会话消息列表", description = "管理端查看任意用户的会话明细") @Parameter(name = "sessionId", description = "会话编号", required = true, example = "1024") diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java new file mode 100644 index 0000000..dd3527a --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java @@ -0,0 +1,20 @@ +package cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message; + +import cn.iocoder.yudao.framework.common.pojo.PageParam; +import io.swagger.v3.oas.annotations.media.Schema; +import jakarta.validation.constraints.NotNull; +import lombok.Data; +import lombok.EqualsAndHashCode; +import lombok.ToString; + +@Schema(description = "管理后台 - AI1 会话消息分页 Request VO") +@Data +@EqualsAndHashCode(callSuper = true) +@ToString(callSuper = true) +public class Ai1SessionMessagePageReqVO extends PageParam { + + @Schema(description = "会话编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024") + @NotNull(message = "会话编号不能为空") + private Long sessionId; + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionPageReqVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionPageReqVO.java index 797b41d..b0924da 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionPageReqVO.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionPageReqVO.java @@ -2,7 +2,6 @@ package cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session; import cn.iocoder.yudao.framework.common.pojo.PageParam; import io.swagger.v3.oas.annotations.media.Schema; -import jakarta.validation.constraints.NotNull; import lombok.Data; import lombok.EqualsAndHashCode; import lombok.ToString; @@ -13,8 +12,7 @@ import lombok.ToString; @ToString(callSuper = true) public class Ai1SessionPageReqVO extends PageParam { - @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") - @NotNull(message = "Agent 编号不能为空") + @Schema(description = "Agent 编号，为空时查询全部 Agent", example = "1") private Long agentId; @Schema(description = "标题，模糊匹配", example = "三体") diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionRespVO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionRespVO.java index 139bcfe..846ce3a 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionRespVO.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/session/Ai1SessionRespVO.java @@ -15,6 +15,9 @@ public class Ai1SessionRespVO { @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") private Long agentId; + @Schema(description = "Agent 名称", example = "客服助手") + private String agentName; + @Schema(description = "用户编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1") private Long userId; diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/agent/Ai1AgentDO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/agent/Ai1AgentDO.java index 822c481..0cd6d90 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/agent/Ai1AgentDO.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/agent/Ai1AgentDO.java @@ -82,6 +82,12 @@ public class Ai1AgentDO extends TenantBaseDO { @TableField(typeHandler = LongListTypeHandler.class) private List skillIds; /** + * 是否开启长期记忆 + * + * 开启后，对话时注入记忆，并由 Job 归档已完成的消息 + */ + private Boolean memoryEnabled; + /** * 状态 * * 枚举 {@link CommonStatusEnum}；开启后才可在「Agent 会话」中使用 diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java new file mode 100644 index 0000000..8b65844 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java @@ -0,0 +1,57 @@ +package cn.iocoder.yudao.module.ai1.dal.dataobject.memory; + +import cn.iocoder.yudao.framework.tenant.core.db.TenantBaseDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; +import com.baomidou.mybatisplus.annotation.KeySequence; +import com.baomidou.mybatisplus.annotation.TableId; +import com.baomidou.mybatisplus.annotation.TableName; +import lombok.*; + +/** + * AI1 长期记忆 DO + * + * 每个用户在每个 Agent 下一条未删除的记忆，由归档 Job 根据已完成的消息增量生成； + * 删除为逻辑删除，保留历史记录便于追溯，之后的归档会根据历史消息生成新的记录 + * + * @author 芋道源码 + */ +@TableName("ai1_memory") +@KeySequence("ai1_memory_seq") // 用于 Oracle、PostgreSQL、Kingbase、DB2、H2 数据库的主键自增。如果是 MySQL 等数据库，可不写。 +@Data +@EqualsAndHashCode(callSuper = true) +@ToString(callSuper = true) +@Builder +@NoArgsConstructor +@AllArgsConstructor +public class Ai1MemoryDO extends TenantBaseDO { + + /** + * 编号 + */ + @TableId + private Long id; + /** + * Agent 编号 + * + * 关联 {@link Ai1AgentDO#getId()} + */ + private Long agentId; + /** + * 用户编号 + * + * 关联 AdminUserDO 的 id 编号 + */ + private Long userId; + /** + * 记忆正文，Markdown 格式，按维度分节 + */ + private String content; + /** + * 已归档的最后一条消息编号，即增量归档的水位线 + * + * 关联 {@link Ai1SessionMessageDO#getId()} + */ + private Long lastMessageId; + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/agent/Ai1AgentMapper.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/agent/Ai1AgentMapper.java index 8a6402d..a1bae0e 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/agent/Ai1AgentMapper.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/agent/Ai1AgentMapper.java @@ -28,7 +28,7 @@ public interface Ai1AgentMapper extends BaseMapperX { default List selectListByStatusOrderById(Integer status) { return selectList(new LambdaQueryWrapperX() - .eq(Ai1AgentDO::getStatus, status) + .eqIfPresent(Ai1AgentDO::getStatus, status) .orderByAsc(Ai1AgentDO::getId)); } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java new file mode 100644 index 0000000..cf3e61c --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java @@ -0,0 +1,32 @@ +package cn.iocoder.yudao.module.ai1.dal.mysql.memory; + +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX; +import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import org.apache.ibatis.annotations.Mapper; + +/** + * AI1 长期记忆 Mapper + * + * @author 芋道源码 + */ +@Mapper +public interface Ai1MemoryMapper extends BaseMapperX { + + default PageResult selectPage(Ai1MemoryPageReqVO reqVO) { + return selectPage(reqVO, new LambdaQueryWrapperX() + .eqIfPresent(Ai1MemoryDO::getAgentId, reqVO.getAgentId()) + .eqIfPresent(Ai1MemoryDO::getUserId, reqVO.getUserId()) + .orderByDesc(Ai1MemoryDO::getId)); + } + + /** + * 未加唯一索引，极端并发下可能存在重复记录，取第一条 + */ + default Ai1MemoryDO selectByAgentIdAndUserId(Long agentId, Long userId) { + return selectFirstOne(Ai1MemoryDO::getAgentId, agentId, Ai1MemoryDO::getUserId, userId); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMapper.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMapper.java index 0b1c88f..666cde7 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMapper.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMapper.java @@ -10,6 +10,8 @@ import org.apache.ibatis.annotations.Mapper; import java.util.Collection; import java.util.List; +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList; + /** * AI1 会话 Mapper * @@ -20,7 +22,7 @@ public interface Ai1SessionMapper extends BaseMapperX { default PageResult selectPage(Ai1SessionPageReqVO reqVO) { return selectPage(reqVO, new LambdaQueryWrapperX() - .eq(Ai1SessionDO::getAgentId, reqVO.getAgentId()) + .eqIfPresent(Ai1SessionDO::getAgentId, reqVO.getAgentId()) .likeIfPresent(Ai1SessionDO::getTitle, reqVO.getTitle()) .eqIfPresent(Ai1SessionDO::getUserId, reqVO.getUserId()) .orderByDesc(Ai1SessionDO::getId)); @@ -33,6 +35,14 @@ public interface Ai1SessionMapper extends BaseMapperX { .orderByDesc(Ai1SessionDO::getId)); } + default List selectUserIdListByAgentId(Long agentId) { + return convertList(selectList(new LambdaQueryWrapperX() + .eq(Ai1SessionDO::getAgentId, agentId) + .select(Ai1SessionDO::getUserId) + .groupBy(Ai1SessionDO::getUserId) + .orderByAsc(Ai1SessionDO::getUserId)), Ai1SessionDO::getUserId); + } + default List selectListByAgentIds(Collection agentIds) { return selectList(new LambdaQueryWrapperX().in(Ai1SessionDO::getAgentId, agentIds)); } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java index 6999c76..ed6b1b4 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java @@ -1,14 +1,18 @@ package cn.iocoder.yudao.module.ai1.dal.mysql.session; +import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX; import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX; import cn.iocoder.yudao.framework.mybatis.core.query.QueryWrapperX; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import org.apache.ibatis.annotations.Mapper; import java.util.Collection; import java.util.List; +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList; + /** * AI1 会话消息 Mapper * @@ -17,6 +21,12 @@ import java.util.List; @Mapper public interface Ai1SessionMessageMapper extends BaseMapperX { + default PageResult selectPage(Ai1SessionMessagePageReqVO reqVO) { + return selectPage(reqVO, new LambdaQueryWrapperX() + .eq(Ai1SessionMessageDO::getSessionId, reqVO.getSessionId()) + .orderByDesc(Ai1SessionMessageDO::getId)); + } + default List selectListBySessionId(Long sessionId) { return selectList(new LambdaQueryWrapperX() .eq(Ai1SessionMessageDO::getSessionId, sessionId) @@ -24,6 +34,18 @@ public interface Ai1SessionMessageMapper extends BaseMapperX selectIdListBySessionIdAndRoleAndStatus(Long sessionId, String role, Integer status) { + return convertList(selectList(new LambdaQueryWrapperX() + .select(Ai1SessionMessageDO::getId) + .eq(Ai1SessionMessageDO::getSessionId, sessionId) + .eq(Ai1SessionMessageDO::getRole, role) + .eq(Ai1SessionMessageDO::getStatus, status) + .orderByAsc(Ai1SessionMessageDO::getId)), Ai1SessionMessageDO::getId); + } + + /** * 按编号倒序，取指定消息之前最近的 limit 条 */ default List selectListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit) { @@ -34,6 +56,32 @@ public interface Ai1SessionMessageMapper extends BaseMapperX sessionIds, Integer status, Long id) { + return selectOne(new QueryWrapperX() + .limitN(1) + .in("session_id", sessionIds) + .eq("status", status) + .gt("id", id) + .orderByAsc("id")); + } + + /** + * 获得多个会话中，编号在 (minId, maxId) 之间的指定状态消息，按编号升序取前 limit 条；maxId 为空时不限上界 + */ + default List selectListBySessionIdsAndStatusAndIdBetween(Collection sessionIds, Integer status, + Long minId, Long maxId, Integer limit) { + return selectList(new QueryWrapperX() + .limitN(limit) + .in("session_id", sessionIds) + .eq("status", status) + .gt("id", minId) + .lt(maxId != null, "id", maxId) + .orderByAsc("id")); + } + default Long selectCountBySessionId(Long sessionId) { return selectCount(Ai1SessionMessageDO::getSessionId, sessionId); } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java index 8a1c01d..3bbd84f 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java @@ -39,4 +39,31 @@ public interface Ai1RedisKeyConstants { */ String SESSION_RESULT_STREAM = "ai1:session:result:%d"; + /** + * 会话停止意图，按助手消息编号隔离；覆盖排队中、其他节点执行中的任务 + *

+ * KEY 格式：ai1:session:stop:{messageId} + * VALUE 数据类型：String，过期时间与结果流一致 + */ + String SESSION_STOP = "ai1:session:stop:%d"; + + /** + * 会话停止通知，消息内容为助手消息编号，用于即时唤醒执行节点 + */ + String SESSION_STOP_TOPIC = "ai1:session:stop"; + + /** + * 长期记忆归档锁，按租户互斥；使用 Redisson 看门狗续期 + *

+ * KEY 格式：ai1:memory:archive-lock:{tenantId} + */ + String MEMORY_ARCHIVE_LOCK = "ai1:memory:archive-lock:%s"; + + /** + * 长期记忆快照锁：消息落库事务与归档读取互斥，避免低编号消息晚提交后落在水位线之前 + *

+ * KEY 格式：ai1:memory:snapshot-lock:{agentId}:{userId}；Agent 编号全局唯一，KEY 不带 tenantId + */ + String MEMORY_SNAPSHOT_LOCK = "ai1:memory:snapshot-lock:%d:%d"; + } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java new file mode 100644 index 0000000..03eff64 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java @@ -0,0 +1,73 @@ +package cn.iocoder.yudao.module.ai1.dal.redis.memory; + +import jakarta.annotation.Resource; +import org.redisson.api.RLock; +import org.redisson.api.RedissonClient; +import org.springframework.stereotype.Repository; + +import java.util.function.Supplier; + +import static cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants.MEMORY_ARCHIVE_LOCK; +import static cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants.MEMORY_SNAPSHOT_LOCK; + +/** + * AI1 长期记忆的锁 Redis DAO + * + * 不设置租期，由 Redisson 看门狗续期 + * + * @author 芋道源码 + */ +@Repository +public class Ai1MemoryLockRedisDAO { + + @Resource + private RedissonClient redissonClient; + + /** + * 按租户尝试获取归档锁并执行；其他节点归档中时不等待、不执行 + * + * @param tenantId 租户编号 + * @param supplier 归档逻辑 + * @return 执行结果；未获取到锁时返回 null + */ + public V tryLockArchive(Long tenantId, Supplier supplier) { + RLock lock = redissonClient.getLock(formatArchiveKey(tenantId)); + if (!lock.tryLock()) { + return null; + } + try { + return supplier.get(); + } finally { + if (lock.isHeldByCurrentThread()) { + lock.unlock(); + } + } + } + + /** + * 获取用户在 Agent 下的快照锁并执行：消息落库事务与归档读取新消息互斥 + * + * @param agentId Agent 编号 + * @param userId 用户编号 + * @param supplier 加锁执行的逻辑 + * @return 执行结果 + */ + public V lockSnapshot(Long agentId, Long userId, Supplier supplier) { + RLock lock = redissonClient.getLock(formatSnapshotKey(agentId, userId)); + lock.lock(); + try { + return supplier.get(); + } finally { + lock.unlock(); + } + } + + private static String formatArchiveKey(Long tenantId) { + return String.format(MEMORY_ARCHIVE_LOCK, tenantId); + } + + private static String formatSnapshotKey(Long agentId, Long userId) { + return String.format(MEMORY_SNAPSHOT_LOCK, agentId, userId); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java index 1bb1b99..69a7b16 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java @@ -77,4 +77,7 @@ public interface Ai1ErrorCodeConstants { // ========== 会话消息 1-041-009-000 ========== ErrorCode SESSION_MESSAGE_NOT_EXISTS = new ErrorCode(1_041_009_000, "会话消息不存在"); + // ========== 长期记忆 1-041-010-000 ========== + ErrorCode MEMORY_NOT_EXISTS = new ErrorCode(1_041_010_000, "长期记忆不存在"); + } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java index d36d982..6e970be 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java @@ -27,6 +27,11 @@ public class YudaoAi1Properties { */ private Skill skill = new Skill(); + /** + * 长期记忆 + */ + private Memory memory = new Memory(); + @Data public static class Milvus { @@ -143,4 +148,19 @@ public class YudaoAi1Properties { } + @Data + public static class Memory { + + /** + * 默认记忆正文最大长度（字符） + */ + public static final int DEFAULT_MAX_LENGTH = 2000; + + /** + * 记忆正文最大长度（字符），超出时截断 + */ + private Integer maxLength = DEFAULT_MAX_LENGTH; + + } + } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java index ef1dace..4657712 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java @@ -90,15 +90,15 @@ public class Ai1LlmChatTool { // 2.1 思考过程：Spring AI 1.x 的 reasoning_content 是增量分片，直接转发（重复文本也可能是有效增量） Object reasoningValue = generation.getOutput().getMetadata().get(METADATA_REASONING); String delta = reasoningValue instanceof String value ? value : null; - if (StrUtil.isNotBlank(delta)) { + if (StrUtil.isNotEmpty(delta)) { thinkingText.append(delta); if (onThinking != null) { onThinking.accept(delta); } } - // 2.2 回复内容 + // 2.2 回复内容：只跳过空串，空格、换行分片也是有效内容 String chunk = generation.getOutput().getText(); - if (StrUtil.isNotBlank(chunk)) { + if (StrUtil.isNotEmpty(chunk)) { contentText.append(chunk); if (onContent != null) { onContent.accept(chunk); diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java new file mode 100644 index 0000000..ae9da92 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java @@ -0,0 +1,193 @@ +package cn.iocoder.yudao.module.ai1.harness.memory; + +import cn.hutool.core.collection.CollUtil; +import cn.hutool.core.util.ObjUtil; +import cn.hutool.core.util.StrUtil; +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; +import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO; +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; +import cn.iocoder.yudao.module.ai1.harness.session.Ai1SessionStreamTool; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService; +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; +import jakarta.annotation.Resource; +import lombok.extern.slf4j.Slf4j; +import org.springframework.stereotype.Component; + +import java.util.Collections; +import java.util.List; + +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList; + +/** + * AI1 长期记忆归档工具（harness 运行时） + * + * 【归档流程】由 Ai1MemoryArchiveJob 按租户调用 {@link #archive} + * 1. 竞争：按租户 tryLock 归档锁，抢不到直接跳过 + * 2. 检测：遍历开启且开启记忆的 Agent，对其下每个用户，读取水位线 lastMessageId 之后的已完成消息 + * 3. 生成：「历史记忆 + 新消息」经模型生成最新记忆，写入并推进水位线 + * + * 【关键约定】 + * 1. 新消息只取「安全前缀」：截止到水位线之后第一条生成中的消息，避免水位线越过尚未完成的回复 + * 2. 超时且无 worker 执行的孤儿生成先按中断收尾，避免永久阻塞安全前缀 + * 3. 单次最多归档 {@link #BATCH_SIZE} 条，剩余的由下次执行继续 + * 4. 写入前校验：记忆被人工修改、删除，或 Agent 已删除时放弃本次结果 + * 5. 锁顺序固定为「归档锁 → 消息执行锁（tryLock）」，快照锁只用于读取新消息，持有期间不收尾、不调用模型 + * + * @author 芋道源码 + */ +@Component +@Slf4j +public class Ai1MemoryArchiveTool { + + /** + * 单个用户单次归档的最大消息数 + */ + static final int BATCH_SIZE = 100; + + @Resource + private Ai1MemoryLockRedisDAO memoryLockRedisDAO; + @Resource + private YudaoAi1Properties ai1Properties; + + @Resource + private Ai1AgentService agentService; + @Resource + private Ai1SessionService sessionService; + @Resource + private Ai1SessionMessageService sessionMessageService; + @Resource + private Ai1MemoryService memoryService; + + @Resource + private Ai1MemoryTool memoryTool; + @Resource + private Ai1SessionStreamTool sessionStreamTool; + + /** + * 归档当前租户下全部开启记忆的 Agent + * + * @return 归档成功的用户数 + */ + public int archive() { + // 1. 按租户竞争归档锁，其他节点归档中时直接跳过 + Integer count = memoryLockRedisDAO.tryLockArchive(TenantContextHolder.getTenantId(), this::archiveAll); + if (count == null) { + log.info("[archive][租户({}) 归档进行中，跳过本次]", TenantContextHolder.getTenantId()); + return 0; + } + return count; + } + + private int archiveAll() { + // 2. 逐个 Agent、逐个用户归档；单个用户失败不影响其他用户 + int count = 0; + int maxLength = ai1Properties.getMemory().getMaxLength(); + for (Ai1AgentDO agent : agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())) { + if (!Boolean.TRUE.equals(agent.getMemoryEnabled())) { + continue; + } + for (Long userId : sessionService.getUserIdListByAgentId(agent.getId())) { + if (Thread.currentThread().isInterrupted()) { + return count; + } + try { + if (archiveOne(agent, userId, maxLength)) { + count++; + } + } catch (Exception e) { + log.warn("[archive][Agent({}) 用户({}) 归档失败]", agent.getId(), userId, e); + } + } + } + return count; + } + + /** + * 归档单个用户：读取新消息 → 生成记忆 → 校验后写入并推进水位线 + * + * @return 是否写入成功 + */ + boolean archiveOne(Ai1AgentDO agent, Long userId, int maxLength) { + // 1.1 读取水位线，并收尾水位线之后超时的孤儿生成 + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(agent.getId(), userId); + Long lastMessageId = memory != null ? ObjUtil.defaultIfNull(memory.getLastMessageId(), 0L) : 0L; + failStaleMessages(agent.getId(), userId, lastMessageId); + // 1.2 读取水位线之后的新消息，无则跳过 + List messages = getNewMessageList(agent.getId(), userId, lastMessageId); + if (CollUtil.isEmpty(messages)) { + return false; + } + + // 2. 生成最新记忆 + String content = memoryTool.generate(agent, memory != null ? memory.getContent() : null, messages, maxLength); + if (StrUtil.isBlank(content)) { + return false; + } + + // 3.1 首次归档（含删除后重新生成）：Agent 已删除，或记忆已存在时放弃，由下次执行基于新记忆继续 + Long newLastMessageId = CollUtil.getLast(messages).getId(); + if (memory == null) { + if (!memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(agent.getId()).userId(userId) + .content(content).lastMessageId(newLastMessageId).build())) { + log.info("[archiveOne][Agent({}) 用户({}) Agent 已删除或记忆已存在，放弃本次结果]", agent.getId(), userId); + return false; + } + return true; + } + // 3.2 再次归档：记忆被人工修改或删除时放弃 + if (!memoryService.updateMemoryIfUnchanged(memory, content, newLastMessageId)) { + log.info("[archiveOne][Agent({}) 用户({}) 记忆已变更，放弃本次结果]", agent.getId(), userId); + return false; + } + return true; + } + + /** + * 收尾水位线之后超时的孤儿生成，直到遇到仍在执行或未超时的生成 + * + * 在快照锁外执行：收尾需抢消息执行锁并读写结果流，避免持有快照锁期间阻塞发送消息 + */ + private void failStaleMessages(Long agentId, Long userId, Long lastMessageId) { + List sessions = sessionService.getSessionListByAgentIdAndUserId(agentId, userId); + if (CollUtil.isEmpty(sessions)) { + return; + } + List sessionIds = convertList(sessions, Ai1SessionDO::getId); + for (int i = 0; i < BATCH_SIZE; i++) { + Ai1SessionMessageDO generatingMessage = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt( + sessionIds, Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId); + if (generatingMessage == null || !sessionStreamTool.failStaleMessage(generatingMessage)) { + return; + } + } + } + + /** + * 获得水位线之后的新消息：截止到第一条生成中的消息，按编号升序，最多 {@link #BATCH_SIZE} 条 + * + * 持有与消息创建相同的快照锁读取：创建在锁内提交，读取时不存在编号更小、尚未提交的消息 + */ + private List getNewMessageList(Long agentId, Long userId, Long lastMessageId) { + return memoryLockRedisDAO.lockSnapshot(agentId, userId, () -> { + List sessions = sessionService.getSessionListByAgentIdAndUserId(agentId, userId); + if (CollUtil.isEmpty(sessions)) { + return Collections.emptyList(); + } + List sessionIds = convertList(sessions, Ai1SessionDO::getId); + Ai1SessionMessageDO generatingMessage = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt( + sessionIds, Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId); + return sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), lastMessageId, + generatingMessage != null ? generatingMessage.getId() : null, BATCH_SIZE); + }); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java new file mode 100644 index 0000000..b4a352c --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java @@ -0,0 +1,162 @@ +package cn.iocoder.yudao.module.ai1.harness.memory; + +import cn.hutool.core.collection.CollUtil; +import cn.hutool.core.util.StrUtil; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; +import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; +import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; +import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService; +import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO; +import cn.iocoder.yudao.module.ai1.util.Ai1Utils; +import jakarta.annotation.Resource; +import org.springframework.stereotype.Component; + +import java.util.Collections; +import java.util.List; + +import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception; +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MODEL_TYPE_NOT_CHAT; + +/** + * AI1 长期记忆工具（harness 运行时） + * + * 1. {@link #loadMemory}：读取用户在 Agent 下的记忆 + * 2. {@link #mergePrompt}：把记忆追加到系统指令，供对话注入 + * 3. {@link #generate}：以「历史记忆 + 最新对话」为输入，经模型按固定维度生成最新记忆 + * + * 记忆的持久化由归档 {@link Ai1MemoryArchiveTool} 完成 + * + * @author 芋道源码 + */ +@Component +public class Ai1MemoryTool { + + /** + * 记忆维度小节标题 + */ + private static final List MEMORY_SECTIONS = List.of("用户画像", "偏好习惯", "关键事实", "目标与意图", "近期进展"); + + /** + * 记忆维度与生成约束（系统指令，{maxLen} 为长度上限占位符） + */ + private static final String MEMORY_SYSTEM_PROMPT = """ + 你是长期记忆归档助手。请基于【历史记忆】与【最新对话】，按下列维度抽取并更新用户长期记忆。 + + 输出格式（务必严格遵守）： + 1. 只输出 Markdown 记忆正文，不要任何解释、前言或代码块围栏； + 2. 每个维度标题独占一行，形如“## 维度名”（井号后有且仅有一个空格）； + 3. 每个要点独占一行，以“- ”开头； + 4. 标题与要点之间、要点与要点之间都必须换行，禁止把标题与内容写在同一行； + 5. 不同维度之间空一行。 + + 维度（无内容的维度整节省略）： + ## 用户画像 + - 用户的身份、称呼、职业、所在地等基本画像 + ## 偏好习惯 + - 用户的沟通偏好、回答风格、常用语言与工具等 + ## 关键事实 + - 用户明确告知、需要长期记住的事实性信息 + ## 目标与意图 + - 用户的核心诉求、长期目标与待办事项 + ## 近期进展 + - 最近对话中已推进的事项、结论与下一步 + + 要求：合并历史记忆与新对话信息，删除重复与已过时内容，保持精炼；总长度不超过 {maxLen} 字。"""; + + private static final String PLACEHOLDER_MAX_LEN = "{maxLen}"; + + /** + * 合并进系统指令时，记忆正文前的引导语 + */ + private static final String MEMORY_PROMPT_PREFIX = "\n\n以下是该用户的历史长期记忆（Markdown），请在回答时结合这些信息；" + + "若与用户当前表述冲突，以当前表述为准：\n"; + + @Resource + private Ai1MemoryService memoryService; + @Resource + private Ai1ModelService modelService; + + @Resource + private Ai1LlmChatTool llmChatTool; + + /** + * 读取用户在 Agent 下的长期记忆 + * + * @param agentId Agent 编号 + * @param userId 用户编号 + * @return 记忆正文；不存在时返回 null + */ + public String loadMemory(Long agentId, Long userId) { + if (agentId == null || userId == null) { + return null; + } + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(agentId, userId); + return memory != null ? memory.getContent() : null; + } + + /** + * 把长期记忆合并进系统指令；记忆为空时原样返回 + * + * @param basePrompt 系统指令 + * @param memory 记忆正文 + * @return 合并后的系统指令 + */ + public String mergePrompt(String basePrompt, String memory) { + String prompt = StrUtil.nullToEmpty(basePrompt); + if (StrUtil.isBlank(memory)) { + return prompt; + } + return prompt + MEMORY_PROMPT_PREFIX + memory.trim(); + } + + /** + * 经模型生成最新记忆：不装配历史、工具、RAG，也不注入已有记忆 + * + * @param agent 记忆所属 Agent，使用其对话模型 + * @param oldMemory 历史记忆，可为空 + * @param newMessages 自上次归档以来的新消息，按编号升序 + * @param maxLength 记忆长度上限（字符） + * @return 最新记忆；模型无输出时返回空 + */ + public String generate(Ai1AgentDO agent, String oldMemory, List newMessages, int maxLength) { + // 1. 校验对话模型 + Ai1ModelRespBO model = modelService.getModelRespBO(agent.getProviderId(), agent.getModelId()); + if (!Ai1ModelTypeEnum.isChat(model.getModelType())) { + throw exception(MODEL_TYPE_NOT_CHAT); + } + + // 2. 单轮会话，聚合为完整文本 + int limit = maxLength > 0 ? maxLength : YudaoAi1Properties.Memory.DEFAULT_MAX_LENGTH; + String systemPrompt = MEMORY_SYSTEM_PROMPT.replace(PLACEHOLDER_MAX_LEN, String.valueOf(limit)); + Ai1LlmChatTool.ChatText chatText = llmChatTool.chat(model, systemPrompt, Collections.emptyList(), + buildUserPrompt(oldMemory, newMessages), Collections.emptyList(), Collections.emptyList(), + null, null, null); + + // 3. 规范化格式并截断 + return Ai1Utils.limitLength(Ai1Utils.normalizeMemoryMarkdown(chatText != null ? chatText.getContent() : null, + MEMORY_SECTIONS), limit); + } + + /** + * 组装归档的用户提示：历史记忆 + 最新对话 + */ + private static String buildUserPrompt(String oldMemory, List newMessages) { + StringBuilder prompt = new StringBuilder(); + prompt.append("【历史记忆】\n"); + prompt.append(StrUtil.isBlank(oldMemory) ? "（暂无）" : oldMemory.trim()).append("\n\n"); + prompt.append("【最新对话】\n"); + if (CollUtil.isNotEmpty(newMessages)) { + for (Ai1SessionMessageDO message : newMessages) { + String role = Ai1SessionMessageRoleEnum.isAssistant(message.getRole()) ? "助手" : "用户"; + prompt.append(role).append("：").append(StrUtil.nullToEmpty(message.getContent())).append("\n"); + } + } + return prompt.toString(); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamTool.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamTool.java index e190cf5..6ee68a6 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamTool.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamTool.java @@ -6,6 +6,7 @@ import cn.hutool.core.map.MapUtil; import cn.hutool.core.thread.ExecutorBuilder; import cn.hutool.core.thread.ThreadUtil; import cn.hutool.core.util.NumberUtil; +import cn.hutool.core.util.ObjUtil; import cn.hutool.core.util.StrUtil; import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; import cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil; @@ -13,6 +14,7 @@ import cn.iocoder.yudao.framework.common.util.json.JsonUtils; import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder; import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils; import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.knowledge.Ai1KnowledgeBaseDO; import cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants; @@ -22,6 +24,8 @@ import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool; +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker; import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory; import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; @@ -33,12 +37,15 @@ import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool; import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory; import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool; import jakarta.annotation.Resource; +import lombok.Getter; import lombok.extern.slf4j.Slf4j; import org.springframework.ai.chat.client.advisor.api.Advisor; import org.springframework.context.SmartLifecycle; import org.springframework.context.annotation.Lazy; +import org.redisson.api.RBucket; import org.redisson.api.RStream; import org.redisson.api.RLock; +import org.redisson.api.RTopic; import org.redisson.api.RedissonClient; import org.redisson.api.stream.*; import org.redisson.client.codec.StringCodec; @@ -46,10 +53,14 @@ import org.springframework.stereotype.Component; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; import java.time.Duration; +import java.time.LocalDateTime; import java.util.*; +import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; @@ -62,6 +73,7 @@ import java.util.function.Consumer; * 2. 生成（worker 侧，{@link #start} 拉起的消费线程）：{@link #consumeLoop} → {@link #handleTask} → {@link #generate}； * 增量写入结果流，结束后回填助手消息并 XACK * 3. 下发与续传（连接侧）：{@link #streamResult} 从结果流 XREAD 转发为 SSE；携带 lastEventId 时从断点之后重放，不重新生成 + * 4. 停止（请求侧）：{@link #cancel} 按助手消息写入停止意图并广播，执行节点中断生成；排队中的任务直接收尾 * * 【关键约定】 * 1. 助手消息编号同时作为结果流标识；生成与连接解耦，任意节点都可以转发、续传，worker 可独立扩容 @@ -86,6 +98,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle { * 生成 worker 线程名前缀，完整线程名为「前缀 + 消费者名」 */ private static final String WORKER_THREAD_NAME_PREFIX = "ai1-session-"; + /** + * 停止意图检查线程名 + */ + private static final String STOP_THREAD_NAME = "ai1-session-stop-"; + /** + * Shell 进程回收线程名前缀，完整线程名为「前缀 + worker 线程名」 + */ + private static final String REAPER_THREAD_NAME_PREFIX = "ai1-session-reap-"; /** * 任务队列最大长度（近似裁剪） @@ -120,6 +140,10 @@ public class Ai1SessionStreamTool implements SmartLifecycle { */ private static final long WORKER_JOIN_TIMEOUT_MILLIS = 3000; /** + * worker 停止时，等待全部 Shell 进程回收线程退出的最长时长，单位：毫秒；单个进程回收最多等待 5s + */ + private static final long REAPER_JOIN_TIMEOUT_MILLIS = 6000; + /** * 结果流连续读取失败的最大次数，达到后中断转发 */ private static final int READ_ERROR_MAX_COUNT = 3; @@ -129,6 +153,11 @@ public class Ai1SessionStreamTool implements SmartLifecycle { private static final long READ_ERROR_BACKOFF_MILLIS = 500; /** + * 停止意图的检查间隔，单位：毫秒；兜底停止通知在订阅断开期间丢失 + */ + private static final long STOP_CHECK_INTERVAL_MILLIS = 1000; + + /** * 结果流丢失判定：转发到连接超时前该时长，结果流仍不存在时，判定任务丢失（预留两个阻塞读窗口，保证 error 能在连接超时前送达） */ private static final long LOST_CHECK_ADVANCE_MILLIS = READ_BLOCK_MILLIS * 2; @@ -143,6 +172,21 @@ public class Ai1SessionStreamTool implements SmartLifecycle { */ static final String INTERRUPTED_MESSAGE = "生成任务已中断，请重新发送"; + /** + * 人工终止、生成失败时追加到正文的提示；失败不拼接原始异常 + */ + static final String TERMINATED_NOTICE = "本条回复已被人工终止"; + static final String FAILED_NOTICE = "本条回复生成失败"; + /** + * 正文与提示之间的分隔 + */ + private static final String NOTICE_PREFIX = "\n\n> "; + + /** + * SSE 转发异常的提示；只是连接异常，后台生成不受影响，可刷新后续传 + */ + static final String TRANSPORT_ERROR_MESSAGE = "会话连接异常，请刷新后重试"; + private static final String FIELD_MESSAGE_ID = "messageId"; private static final String FIELD_AGENT_ID = "agentId"; private static final String FIELD_SESSION_ID = "sessionId"; @@ -173,6 +217,8 @@ public class Ai1SessionStreamTool implements SmartLifecycle { private Ai1LlmChatTool llmChatTool; @Resource private Ai1RagTool ragTool; + @Resource + private Ai1MemoryTool memoryTool; @Resource private Ai1McpToolFactory mcpToolFactory; @@ -184,6 +230,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle { */ private final Map resultExpireTimes = new ConcurrentHashMap<>(); /** + * 本节点执行中的生成：助手消息编号 → 运行上下文，供停止时中断生成线程、回收 Shell 进程 + */ + private final Map runContexts = new ConcurrentHashMap<>(); + /** + * 未结束的 Shell 进程回收线程：生成结束后可能仍在回收，停止时统一等待；线程结束时自行移除 + */ + private final Set reapers = ConcurrentHashMap.newKeySet(); + /** * 实例标识，消费者名按「实例 + 序号」隔离 */ private final String instanceId = UUID.randomUUID().toString().substring(0, 8); @@ -199,6 +253,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle { * SSE 转发线程池：单个转发阻塞于 XREAD 直到流结束，用有界池控制并发 */ private ExecutorService sseExecutor; + /** + * 停止通知的订阅编号 + */ + private Integer stopListenerId; + /** + * 停止意图的定时检查线程 + */ + private ScheduledExecutorService stopChecker; // ==================== 生命周期 ==================== @@ -224,7 +286,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle { // 2. 初始化任务消费组 initTaskGroup(); - // 3. 拉起 worker 消费线程，消费者名按「实例 + 序号」隔离 + // 3. 订阅停止通知，即时中断本节点的生成；通知不持久，另定时检查停止意图兜底 + // 中断交给检查线程执行，避免 Redisson 订阅线程等待增量回调 + stopChecker = Executors.newSingleThreadScheduledExecutor(ThreadUtil.newNamedThreadFactory(STOP_THREAD_NAME, true)); + stopChecker.scheduleWithFixedDelay(this::checkStopRequests, + STOP_CHECK_INTERVAL_MILLIS, STOP_CHECK_INTERVAL_MILLIS, TimeUnit.MILLISECONDS); + stopListenerId = getStopTopic().addListener(String.class, + (channel, messageId) -> stopChecker.execute(() -> cancelRunContext(Convert.toLong(messageId)))); + + // 4. 拉起 worker 消费线程，消费者名按「实例 + 序号」隔离 running = true; int workerCount = ai1Properties.getSession().getWorker().getCount(); for (int i = 0; i < workerCount; i++) { @@ -237,11 +307,33 @@ public class Ai1SessionStreamTool implements SmartLifecycle { } /** - * 停止：置停止标志，等待各消费线程退出（每个最多 3s），最后关闭 SSE 转发线程池 + * 停止：置停止标志，取消停止订阅，中断本节点执行中的生成，等待各消费线程退出（每个最多 3s）、 + * Shell 进程回收（共最多 6s），最后关闭 SSE 转发线程池 + * + * 中断的生成保留部分内容并标记失败，不会被其他节点认领后重新生成 */ @Override public void stop() { + // 1. 先置停止标志，再在 runContexts 锁内取快照：与 registerRunContext 互斥， + // 之后登记的上下文由登记方自行中断，不会漏掉阻塞读取刚返回、尚未登记的任务 running = false; + if (stopListenerId != null) { + try { + getStopTopic().removeListener(stopListenerId); + } catch (Exception e) { + log.warn("[stop][停止通知取消订阅失败]", e); + } + } + if (stopChecker != null) { + stopChecker.shutdownNow(); + } + List contexts; + synchronized (runContexts) { + contexts = new ArrayList<>(runContexts.values()); + } + // 中断只启动各自的回收线程，单个进程回收慢不影响其他上下文 + contexts.forEach(context -> context.cancel(INTERRUPTED_MESSAGE)); + // 2. 等待消费线程退出 for (Thread worker : workers) { try { worker.join(WORKER_JOIN_TIMEOUT_MILLIS); @@ -250,6 +342,19 @@ public class Ai1SessionStreamTool implements SmartLifecycle { } } workers.clear(); + // 3. 等待全部未结束的 Shell 进程回收，包括此前人工停止、生成已结束的回收；各回收线程并行，共用一个截止时刻 + long reaperDeadline = System.currentTimeMillis() + REAPER_JOIN_TIMEOUT_MILLIS; + try { + for (Thread reaper : reapers) { + long remaining = reaperDeadline - System.currentTimeMillis(); + if (remaining <= 0) { + break; + } + reaper.join(remaining); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } if (sseExecutor != null) { sseExecutor.shutdown(); } @@ -268,15 +373,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle { try { getTaskStream().createGroup(StreamCreateGroupArgs.name(Ai1RedisKeyConstants.SESSION_TASK_GROUP).id(StreamMessageId.ALL).makeStream()); } catch (Exception e) { - if (!isBusyGroup(e)) { + if (!isRedisError(e, "BUSYGROUP")) { log.warn("[initTaskGroup][Chat 任务消费组初始化失败]", e); } } } - private static boolean isBusyGroup(Throwable e) { + private static boolean isRedisError(Throwable e, String code) { for (Throwable cause = e; cause != null; cause = cause.getCause()) { - if (StrUtil.contains(cause.getMessage(), "BUSYGROUP")) { + if (StrUtil.contains(cause.getMessage(), code)) { return true; } } @@ -331,8 +436,9 @@ public class Ai1SessionStreamTool implements SmartLifecycle { emitter.complete(); } } catch (Exception e) { + // 只是转发异常，不改动消息状态，也不按生成失败提示 log.warn("[open][助手消息({}) 会话流转发异常]", messageId, e); - sendErrorAndComplete(emitter, e.getMessage()); + sendErrorAndComplete(emitter, TRANSPORT_ERROR_MESSAGE); } }); } catch (RejectedExecutionException e) { @@ -354,6 +460,151 @@ public class Ai1SessionStreamTool implements SmartLifecycle { return emitter; } + // ==================== 停止 ==================== + + /** + * 停止会话中生成中的助手消息 + * + * 按助手消息写入停止意图并广播，覆盖其他节点执行中、排队中的任务 + * + * @param sessionId 会话编号 + */ + public void cancel(Long sessionId) { + // 1. 按请求时刻定位生成中的助手消息；之后发送的新一轮不受影响 + List messageIds = sessionMessageService.getGeneratingAssistantMessageIdList(sessionId); + Duration ttl = Duration.ofSeconds(ai1Properties.getSession().getStream().getTtlSeconds()); + for (Long messageId : messageIds) { + // 2.1 写入停止意图：worker 执行前、执行中都会检查，与结果流的过期时间一致 + getStopBucket(messageId).set("1", ttl); + // 2.2 广播停止通知，执行节点即时中断 + getStopTopic().publish(String.valueOf(messageId)); + // 2.3 没有 worker 执行时（排队中，或 worker 已退出），直接收尾 + cancelQueuedTask(messageId); + } + } + + /** + * 收尾无人执行的任务：抢到执行锁说明没有 worker 在生成，保留已有结果并标记终止 + * + * 任务记录仍在队列中，之后由 worker 按「已回填」补齐终态并确认，不再调用模型 + */ + private void cancelQueuedTask(Long messageId) { + RLock lock = redissonClient.getLock(String.format(Ai1RedisKeyConstants.SESSION_TASK_LOCK, messageId)); + if (!lock.tryLock()) { + return; + } + try { + Ai1SessionMessageDO message = sessionMessageService.getSessionMessage(messageId); + if (message == null || !Ai1SessionMessageStatusEnum.isGenerating(message.getStatus())) { + return; + } + saveAssistantMessage(messageId, message.getSessionId(), withNotice(readPartialResult(messageId), TERMINATED_NOTICE), + Ai1SessionMessageStatusEnum.FAILED.getStatus()); + appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), ""); + } finally { + if (lock.isHeldByCurrentThread()) { + lock.unlock(); + } + } + } + + /** + * 收尾超时的孤儿生成：任务记录被裁剪或丢失时，助手消息会一直停在生成中，阻塞长期记忆归档 + * + * 超时阈值与超时任务认领一致，仍有任务记录的消息此前已被认领收尾； + * 先判定超时再抢执行锁，避免短暂持锁使 worker 把排队中的新任务误判为在途而跳过 + * + * @param message 生成中的助手消息 + * @return 是否已不在生成中；仍在执行或未超时时返回 false + * @throws IllegalStateException 收尾期间执行锁失效，由归档按单个用户失败处理 + */ + public boolean failStaleMessage(Ai1SessionMessageDO message) { + // 1. 未超时直接保留，不加锁 + if (!isStaleMessage(message)) { + return false; + } + // 2. 抢不到执行锁说明仍有 worker 在生成 + Long messageId = message.getId(); + RLock lock = redissonClient.getLock(String.format(Ai1RedisKeyConstants.SESSION_TASK_LOCK, messageId)); + if (!lock.tryLock()) { + return false; + } + long ownerThreadId = Thread.currentThread().threadId(); + try { + // 3. 锁内重新读取，期间可能已被 worker 完成 + Ai1SessionMessageDO current = sessionMessageService.getSessionMessage(messageId); + if (current == null || !Ai1SessionMessageStatusEnum.isGenerating(current.getStatus())) { + return true; + } + if (!isStaleMessage(current)) { + return false; + } + // 4. 与接管任务的收尾一致：保留部分结果并标记失败，再写入 error、done； + // 每次写入前校验执行权，锁失效后由新持有者处理，不再改动消息或结果流 + log.warn("[failStaleMessage][助手消息({}) 生成超时且无 worker 执行，标记中断]", messageId); + Ai1LlmChatTool.ChatText partialText = readPartialResult(messageId); + checkTaskLock(lock, ownerThreadId); + saveAssistantMessage(messageId, current.getSessionId(), withNotice(partialText, INTERRUPTED_MESSAGE), + Ai1SessionMessageStatusEnum.FAILED.getStatus()); + checkTaskLock(lock, ownerThreadId); + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE); + checkTaskLock(lock, ownerThreadId); + appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), ""); + return true; + } finally { + if (lock.isHeldByThread(ownerThreadId)) { + lock.unlock(); + } + } + } + + private boolean isStaleMessage(Ai1SessionMessageDO message) { + if (!Ai1SessionMessageRoleEnum.isAssistant(message.getRole()) || message.getCreateTime() == null) { + return false; + } + long staleMillis = ai1Properties.getSession().getStream().getTimeoutMs() + RECLAIM_IDLE_BUFFER_MILLIS; + return message.getCreateTime().isBefore(LocalDateTime.now().minus(Duration.ofMillis(staleMillis))); + } + + /** + * 登记运行上下文；节点已停机时立即中断 + * + * 与 {@link #stop} 的快照在同一把锁内：要么被快照中断，要么在此看到停止标志 + */ + private void registerRunContext(Long messageId, RunContext context) { + synchronized (runContexts) { + runContexts.put(messageId, context); + if (!running) { + context.cancel(INTERRUPTED_MESSAGE); + } + } + } + + /** + * 中断本节点执行中的生成；不在本节点执行时忽略 + */ + private void cancelRunContext(Long messageId) { + RunContext context = messageId != null ? runContexts.get(messageId) : null; + if (context != null) { + context.cancel(TERMINATED_NOTICE); + } + } + + /** + * 定时检查本节点执行中的生成是否有停止意图，兜底订阅断开期间丢失的停止通知 + */ + private void checkStopRequests() { + runContexts.keySet().forEach(messageId -> { + try { + if (getStopBucket(messageId).isExists()) { + cancelRunContext(messageId); + } + } catch (Exception e) { + log.warn("[checkStopRequests][助手消息({}) 停止意图检查失败]", messageId, e); + } + }); + } + // ==================== worker 消费 ==================== /** @@ -382,6 +633,10 @@ public class Ai1SessionStreamTool implements SmartLifecycle { return; } log.warn("[consumeLoop][Chat worker({}) 消费异常]", consumer, e); + // 任务队列被删除（如 Redis 清库）后消费组随之丢失，重建后继续消费 + if (isRedisError(e, "NOGROUP")) { + initTaskGroup(); + } ThreadUtil.sleep(CONSUME_ERROR_BACKOFF_MILLIS); } } @@ -477,7 +732,8 @@ public class Ai1SessionStreamTool implements SmartLifecycle { Ai1LlmChatTool.ChatText partialText = readPartialResult(messageId); // 1.3.2 回填部分结果并标记失败；每次写入前校验执行权，防止旧 worker 继续写入 checkTaskLock(lock, workerThreadId); - saveAssistantMessage(messageId, sessionId, partialText, Ai1SessionMessageStatusEnum.FAILED.getStatus()); + saveAssistantMessage(messageId, sessionId, withNotice(partialText, INTERRUPTED_MESSAGE), + Ai1SessionMessageStatusEnum.FAILED.getStatus()); // 1.3.3 写入错误提示和结束事件，续传连接可及时退出 checkTaskLock(lock, workerThreadId); appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE); @@ -493,40 +749,81 @@ public class Ai1SessionStreamTool implements SmartLifecycle { Long agentId = Convert.toLong(fields.get(FIELD_AGENT_ID)); String content = fields.get(FIELD_CONTENT); - // 3. 生成并回填；增量回调可能在模型线程执行，因此按 worker 线程编号校验锁归属 + // 3. 登记运行上下文，供停止时中断生成线程、回收 Shell 进程 + RunContext context = new RunContext(Thread.currentThread(), skillToolFactory.newShellTracker(agentId), reapers); + registerRunContext(messageId, context); + // 增量回调可能在模型线程执行，因此按 worker 线程编号校验锁归属；停止后不再写入增量 StringBuilder contentText = new StringBuilder(); StringBuilder thinkingText = new StringBuilder(); Consumer onThinking = delta -> { - checkTaskLock(lock, workerThreadId); - thinkingText.append(delta); - appendResult(messageId, Ai1SessionStreamEventEnum.THINKING.getEvent(), delta); + synchronized (context) { + context.checkActive(); + checkTaskLock(lock, workerThreadId); + thinkingText.append(delta); + appendResult(messageId, Ai1SessionStreamEventEnum.THINKING.getEvent(), delta); + } }; Consumer onContent = delta -> { - checkTaskLock(lock, workerThreadId); - contentText.append(delta); - appendResult(messageId, Ai1SessionStreamEventEnum.MESSAGE.getEvent(), delta); + synchronized (context) { + context.checkActive(); + checkTaskLock(lock, workerThreadId); + contentText.append(delta); + appendResult(messageId, Ai1SessionStreamEventEnum.MESSAGE.getEvent(), delta); + } }; - // 3.1 标记生成开始，再调用模型并回填完整结果 + + // 4. 标记生成开始，再调用模型 + Ai1LlmChatTool.ChatText chatText = null; + Exception failure = null; + String stopNotice; try { + // 4.1 节点已停机，或排队期间已被停止时，不再调用模型 + context.checkActive(); + if (getStopBucket(messageId).isExists()) { + context.cancel(TERMINATED_NOTICE); + } + context.checkActive(); + // 4.2 标记生成已开始；长思考没有增量时，续传也不会误判结果流丢失 checkTaskLock(lock, workerThreadId); - // 标记生成已开始；长思考没有增量时，续传也不会误判结果流丢失 appendResult(messageId, Ai1SessionStreamEventEnum.PING.getEvent(), ""); - Ai1LlmChatTool.ChatText chatText = generate(messageId, agentId, sessionId, content, onThinking, onContent); - checkTaskLock(lock, workerThreadId); - saveAssistantMessage(messageId, sessionId, chatText, Ai1SessionMessageStatusEnum.SUCCESS.getStatus()); + chatText = generate(context, messageId, agentId, sessionId, content, onThinking, onContent); } catch (Exception e) { - // 失去锁的旧 worker 不得再写结果、终态或 XACK，留给新持有者处理 + failure = e; + } finally { + // 4.3 结束上下文并清除中断标志，之后的停止不再生效 + stopNotice = context.finish(); + runContexts.remove(messageId, context); + context.getShellTracker().cancel(); + } + + // 5. 先回填终态，再下发 error / done；每次写入前校验执行权 + if (stopNotice != null) { + // 5.1 停止：保留已生成的内容并追加提示；人工终止只写 done，节点停机中断另写 error + log.info("[handleTask][助手消息({}) 生成已停止：{}]", messageId, stopNotice); + Ai1LlmChatTool.ChatText partialText = chatText != null ? chatText + : new Ai1LlmChatTool.ChatText(contentText.toString(), thinkingText.toString()); + checkTaskLock(lock, workerThreadId); + saveAssistantMessage(messageId, sessionId, withNotice(partialText, stopNotice), Ai1SessionMessageStatusEnum.FAILED.getStatus()); + if (!TERMINATED_NOTICE.equals(stopNotice)) { + checkTaskLock(lock, workerThreadId); + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), stopNotice); + } + } else if (failure != null) { + // 5.2 失败：失去锁的旧 worker 不得再写结果、终态或 XACK，留给新持有者处理 if (!lock.isHeldByThread(workerThreadId)) { return; } - // 3.2 生成失败时下发错误，保留本轮已产生的增量并回填失败状态 - log.warn("[handleTask][助手消息({}) 生成失败]", messageId, e); - appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), StrUtil.blankToDefault(e.getMessage(), "生成失败")); + log.warn("[handleTask][助手消息({}) 生成失败]", messageId, failure); + saveAssistantMessage(messageId, sessionId, withNotice(new Ai1LlmChatTool.ChatText(contentText.toString(), + thinkingText.toString()), FAILED_NOTICE), Ai1SessionMessageStatusEnum.FAILED.getStatus()); checkTaskLock(lock, workerThreadId); - saveAssistantMessage(messageId, sessionId, new Ai1LlmChatTool.ChatText(contentText.toString(), - thinkingText.toString()), Ai1SessionMessageStatusEnum.FAILED.getStatus()); + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), FAILED_NOTICE); + } else { + // 5.3 成功：回填完整结果 + checkTaskLock(lock, workerThreadId); + saveAssistantMessage(messageId, sessionId, chatText, Ai1SessionMessageStatusEnum.SUCCESS.getStatus()); } - // 4. 回填后写终态并确认；失败时由后续接管清理，不能在 finally 中无条件确认 + // 6. 回填后写终态并确认；失败时由后续接管清理，不能在 finally 中无条件确认 checkTaskLock(lock, workerThreadId); appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), ""); checkTaskLock(lock, workerThreadId); @@ -546,11 +843,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle { if (Ai1SessionStreamEventEnum.DONE.getEvent().equals(lastType)) { return; } - // 2. 失败消息缺少 error 时补写错误提示，末尾已有 error 时不重复写入 + // 2. 失败消息缺少 error 时，按正文中已回填的提示补写；人工终止只写 done，末尾已有 error 时不重复写入 if (Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()) && !Ai1SessionStreamEventEnum.ERROR.getEvent().equals(lastType)) { - checkTaskLock(lock, workerThreadId); - appendResult(message.getId(), Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE); + String notice = getSavedNotice(message.getContent()); + if (!TERMINATED_NOTICE.equals(notice)) { + checkTaskLock(lock, workerThreadId); + appendResult(message.getId(), Ai1SessionStreamEventEnum.ERROR.getEvent(), + ObjUtil.defaultIfNull(notice, INTERRUPTED_MESSAGE)); + } } // 3. 补写 done，使续传连接能够结束；调用方在补齐后再确认任务 checkTaskLock(lock, workerThreadId); @@ -593,12 +894,12 @@ public class Ai1SessionStreamTool implements SmartLifecycle { // ==================== 生成编排 ==================== /** - * 执行一次生成：校验 Agent → 解析模型 → 装配历史、工具、RAG → 流式会话 + * 执行一次生成：校验 Agent → 解析模型 → 装配历史、长期记忆、工具、RAG → 流式会话 * * 历史只取当前助手占位之前最近 historyLimit 条，跳过未完成的助手占位，并移除末尾的当前提问（由 LLM 工具另行追加） */ @SuppressWarnings("SequencedCollectionMethodCanBeUsed") - private Ai1LlmChatTool.ChatText generate(Long messageId, Long agentId, Long sessionId, String content, + private Ai1LlmChatTool.ChatText generate(RunContext context, Long messageId, Long agentId, Long sessionId, String content, Consumer onThinking, Consumer onContent) { // 1. 生成可能发生在其他节点，重新校验 Agent 与模型；后台会话不要求 Agent 已发布 Ai1AgentDO agent = agentService.validateAgentExists(agentId); @@ -631,10 +932,13 @@ public class Ai1SessionStreamTool implements SmartLifecycle { String systemPrompt = StrUtil.isNotBlank(agent.getSystemPrompt()) ? agent.getSystemPrompt().trim() : StrUtil.replace(ai1Properties.getSession().getDefaultSystemPrompt(), YudaoAi1Properties.Session.AGENT_NAME_PLACEHOLDER, agent.getName()); + // 开启长期记忆时，合入用户在该 Agent 下的记忆 + systemPrompt = injectMemory(agent, sessionId, systemPrompt); // 4. 工具（MCP + SKILL）与 RAG Advisor；知识库不可用时降级跳过，并把提示写入思考流 + // Shell 命令经运行上下文的跟踪器执行，停止或生成结束时统一回收 List tools = new ArrayList<>(mcpToolFactory.buildTools(agent)); - tools.addAll(skillToolFactory.buildTools(agent)); + tools.addAll(skillToolFactory.buildTools(agent, context.getShellTracker())); StringBuilder noticeText = new StringBuilder(); List advisors = new ArrayList<>(); for (Ai1KnowledgeBaseDO knowledgeBase : knowledgeBaseService.getKnowledgeBaseList(agent.getKnowledgeBaseIds())) { @@ -652,6 +956,7 @@ public class Ai1SessionStreamTool implements SmartLifecycle { } // 5. 流式会话；降级提示并入最终思考文本，保证刷新后内容一致 + context.checkActive(); Ai1LlmChatTool.ChatText chatText = llmChatTool.chat(model, systemPrompt, histories, content, tools, advisors, sessionId, onThinking, onContent); if (!noticeText.isEmpty()) { @@ -661,6 +966,27 @@ public class Ai1SessionStreamTool implements SmartLifecycle { } /** + * 长期记忆注入 + * + * 查询失败时降级为原系统指令，保证对话可用 + */ + private String injectMemory(Ai1AgentDO agent, Long sessionId, String systemPrompt) { + if (!Boolean.TRUE.equals(agent.getMemoryEnabled())) { + return systemPrompt; + } + try { + Ai1SessionDO session = sessionService.getSession(sessionId); + if (session == null) { + return systemPrompt; + } + return memoryTool.mergePrompt(systemPrompt, memoryTool.loadMemory(agent.getId(), session.getUserId())); + } catch (Exception e) { + log.warn("[injectMemory][Agent({}) 会话({}) 长期记忆注入失败，已降级跳过]", agent.getId(), sessionId, e); + return systemPrompt; + } + } + + /** * 回填助手消息（占位记录）的内容、思考过程与生成状态，并刷新会话活跃时间 */ private void saveAssistantMessage(Long messageId, Long sessionId, Ai1LlmChatTool.ChatText chatText, Integer status) { @@ -830,6 +1156,33 @@ public class Ai1SessionStreamTool implements SmartLifecycle { return redissonClient.getStream(String.format(Ai1RedisKeyConstants.SESSION_RESULT_STREAM, messageId), StringCodec.INSTANCE); } + private RBucket getStopBucket(Long messageId) { + return redissonClient.getBucket(String.format(Ai1RedisKeyConstants.SESSION_STOP, messageId), StringCodec.INSTANCE); + } + + private RTopic getStopTopic() { + return redissonClient.getTopic(Ai1RedisKeyConstants.SESSION_STOP_TOPIC, StringCodec.INSTANCE); + } + + /** + * 正文末尾追加提示 + */ + static Ai1LlmChatTool.ChatText withNotice(Ai1LlmChatTool.ChatText chatText, String notice) { + return new Ai1LlmChatTool.ChatText(StrUtil.nullToEmpty(chatText.getContent()) + NOTICE_PREFIX + notice, chatText.getThinking()); + } + + /** + * 解析正文末尾已回填的提示；没有时返回空 + */ + static String getSavedNotice(String content) { + for (String notice : List.of(TERMINATED_NOTICE, FAILED_NOTICE, INTERRUPTED_MESSAGE)) { + if (StrUtil.endWith(content, NOTICE_PREFIX + notice)) { + return notice; + } + } + return null; + } + /** * 解析结果流条目编号（格式为 毫秒时间戳-序号）；为空或非法时从头读取 */ @@ -844,4 +1197,94 @@ public class Ai1SessionStreamTool implements SmartLifecycle { return new StreamMessageId(Long.parseLong(parts.get(0)), Long.parseLong(parts.get(1))); } + /** + * 生成的运行上下文 + * + * 停止与结束互斥，结束后迟到的停止不再中断线程 + */ + static class RunContext { + + /** + * 生成线程，即 worker 线程 + */ + private final Thread thread; + /** + * Shell 进程跟踪器 + */ + @Getter + private final Ai1SkillShellTracker shellTracker; + /** + * 是否已结束 + */ + private boolean finished; + /** + * 停止提示；为空表示未停止 + */ + private String stopNotice; + + /** + * 未结束的 Shell 进程回收线程，与运行上下文的生命周期无关，由节点停止时统一等待 + */ + private final Set reapers; + + RunContext(Thread thread, Ai1SkillShellTracker shellTracker, Set reapers) { + this.thread = thread; + this.shellTracker = shellTracker; + this.reapers = reapers; + } + + /** + * 停止生成：中断生成线程，并由独立线程回收在册的后台进程；已结束、已停止时忽略 + * + * 单个进程回收最多等待 5s：不在调用方执行，避免阻塞停止检查线程或其他上下文的停止； + * 不依赖 worker 响应中断，模型或工具库吞掉中断时进程也会被回收； + * 回收可能晚于生成结束，结束时 worker 拿不到已被回收线程认领的进程，因此回收线程登记到 reapers，不随上下文移除 + * + * @param notice 停止提示 + */ + synchronized void cancel(String notice) { + if (finished || stopNotice != null) { + return; + } + stopNotice = notice; + thread.interrupt(); + Thread reaper = ThreadUtil.newThread(() -> { + try { + shellTracker.cancel(); + } finally { + reapers.remove(Thread.currentThread()); + } + }, REAPER_THREAD_NAME_PREFIX + thread.getName(), true); + // 先登记再启动，结束时必能移除自身 + reapers.add(reaper); + try { + reaper.start(); + } catch (Throwable e) { + reapers.remove(reaper); + throw e; + } + } + + /** + * 校验生成仍在进行，已停止或已结束时中止后续步骤 + */ + synchronized void checkActive() { + if (finished || stopNotice != null) { + throw new CancellationException(StrUtil.blankToDefault(stopNotice, "生成已结束")); + } + } + + /** + * 结束生成，并清除 worker 线程的中断标志；需在 worker 线程调用 + * + * @return 停止提示；未停止时为空 + */ + synchronized String finish() { + finished = true; + Thread.interrupted(); + return stopNotice; + } + + } + } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java new file mode 100644 index 0000000..0d284db --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java @@ -0,0 +1,74 @@ +package cn.iocoder.yudao.module.ai1.harness.skill; + +import lombok.extern.slf4j.Slf4j; +import org.springaicommunity.agent.common.exec.ExecBackend; +import org.springaicommunity.agent.common.exec.ExecHandle; +import org.springaicommunity.agent.common.exec.ExecResult; +import org.springaicommunity.agent.common.exec.ExecSpec; + +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * AI1 SKILL Shell 进程跟踪器 + * + * 登记本次生成启动的后台进程，生成结束或终止时回收 + * + * @author 芋道源码 + */ +@Slf4j +public class Ai1SkillShellTracker implements ExecBackend { + + private final ExecBackend delegate; + /** + * 在册的后台进程 + */ + private final Set handles = ConcurrentHashMap.newKeySet(); + /** + * 是否已终止 + */ + private volatile boolean cancelled = false; + + public Ai1SkillShellTracker(ExecBackend delegate) { + this.delegate = delegate; + } + + @Override + public ExecResult run(ExecSpec spec) { + return delegate.run(spec); + } + + @Override + public ExecHandle start(ExecSpec spec) { + if (cancelled) { + throw new IllegalStateException("生成已终止"); + } + ExecHandle handle = delegate.start(spec); + handles.add(handle); + // 登记后才被终止时，立即回收；先置标志再遍历的 cancel 与此处必有一方看到该进程 + if (cancelled && handles.remove(handle)) { + handle.kill(); + } + return handle; + } + + /** + * 终止并回收全部在册的后台进程（幂等，可并发调用） + * + * 按「移出集合」认领进程，每个进程只回收一次 + */ + public void cancel() { + cancelled = true; + for (ExecHandle handle : handles) { + if (!handles.remove(handle)) { + continue; + } + try { + handle.kill(); + } catch (Exception e) { + log.warn("[cancel][Shell 进程({}) 回收失败]", handle.id(), e); + } + } + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java index 7f08873..ce169bc 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java @@ -18,6 +18,7 @@ import jakarta.annotation.Resource; import lombok.AllArgsConstructor; import lombok.Data; import lombok.extern.slf4j.Slf4j; +import org.springaicommunity.agent.exec.LocalExecBackend; import org.springaicommunity.agent.tools.*; import org.springframework.ai.tool.ToolCallback; import org.springframework.stereotype.Component; @@ -87,10 +88,11 @@ public class Ai1SkillToolFactory { /** * 构建 Agent 绑定 SKILL 所需的全部工具：SkillsTool（注入 SKILL.md 内容）+ 执行类工具（限定于 Agent 物化根目录） * - * @param agent Agent + * @param agent Agent + * @param tracker Shell 进程跟踪器，Shell 命令经其执行，生成终止或结束时统一回收 * @return 工具列表；未绑定有效 SKILL 时返回空列表 */ - public List buildTools(Ai1AgentDO agent) { + public List buildTools(Ai1AgentDO agent, Ai1SkillShellTracker tracker) { // 1. 同步物化快照；未绑定有效 SKILL 时，不提供任何工具 SkillSnapshot snapshot = sync(agent); if (snapshot == null) { @@ -105,7 +107,7 @@ public class Ai1SkillToolFactory { // 3. 执行类工具：ShellTools（bash）、FileSystemTools（Read/Write/Edit）、Glob、Grep、ListDirectory Path root = snapshot.getRoot(); - tools.add(ShellTools.builder().workingDirectory(root).build()); + tools.add(ShellTools.builder().execBackend(tracker).build()); tools.add(FileSystemTools.builder().allowedDirectory(root).build()); tools.add(GlobTool.builder().workingDirectory(root).allowedDirectory(root).build()); tools.add(GrepTool.builder().workingDirectory(root).allowedDirectory(root).build()); @@ -114,6 +116,16 @@ public class Ai1SkillToolFactory { } /** + * 创建 Agent 的 Shell 进程跟踪器：工作目录限定于 Agent 物化根目录 + * + * @param agentId Agent 编号 + * @return Shell 进程跟踪器 + */ + public Ai1SkillShellTracker newShellTracker(Long agentId) { + return new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(buildAgentRoot(agentId)).build()); + } + + /** * 清理 Agent 物化沙箱：摘除快照缓存并删除本地目录（Agent 删除时调用） * * @param agentId Agent 编号 diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java new file mode 100644 index 0000000..b7dc066 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java @@ -0,0 +1,27 @@ +package cn.iocoder.yudao.module.ai1.job; + +import cn.iocoder.yudao.framework.quartz.core.handler.JobHandler; +import cn.iocoder.yudao.framework.tenant.core.job.TenantJob; +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryArchiveTool; +import jakarta.annotation.Resource; +import org.springframework.stereotype.Component; + +/** + * AI1 长期记忆归档 Job + * + * @author 芋道源码 + */ +@Component +public class Ai1MemoryArchiveJob implements JobHandler { + + @Resource + private Ai1MemoryArchiveTool memoryArchiveTool; + + @Override + @TenantJob + public String execute(String param) { + int count = memoryArchiveTool.archive(); + return String.format("AI1 长期记忆归档：归档 %s 个用户", count); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java index 14533f2..e26e4d8 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java @@ -81,9 +81,9 @@ public interface Ai1AgentService { Ai1AgentDO validateAgentEnabled(Long id); /** - * 获得指定状态的 Agent 列表，用于会话页选择 Agent + * 获得指定状态的 Agent 列表，按编号升序，用于选择 Agent * - * @param status 状态 + * @param status 状态；为空时返回全部状态 * @return Agent 列表 */ List getAgentListByStatus(Integer status); diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImpl.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImpl.java index d23a22e..a946e12 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImpl.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImpl.java @@ -10,6 +10,7 @@ import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; import cn.iocoder.yudao.module.ai1.dal.mysql.agent.Ai1AgentMapper; import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService; import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO; @@ -45,6 +46,9 @@ public class Ai1AgentServiceImpl implements Ai1AgentService { @Resource @Lazy // 延迟加载，避免循环依赖 private Ai1SessionService sessionService; + @Resource + @Lazy // 延迟加载，避免循环依赖 + private Ai1MemoryService memoryService; @Resource private Ai1SkillToolFactory skillToolFactory; @@ -82,9 +86,10 @@ public class Ai1AgentServiceImpl implements Ai1AgentService { // 1. 校验存在 ids.forEach(this::validateAgentExists); - // 2. 删除 Agent，同一事务内级联删除会话与消息 + // 2. 删除 Agent，同一事务内级联删除会话、消息与长期记忆 agentMapper.deleteByIds(ids); sessionService.deleteSessionListByAgentIds(ids); + memoryService.deleteMemoryListByAgentIds(ids); // 3. 清理 SKILL 物化沙箱 ids.forEach(skillToolFactory::evict); diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java new file mode 100644 index 0000000..fc52f98 --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java @@ -0,0 +1,90 @@ +package cn.iocoder.yudao.module.ai1.service.memory; + +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import jakarta.validation.Valid; + +import java.util.Collection; +import java.util.List; + +/** + * AI1 长期记忆 Service 接口 + * + * @author 芋道源码 + */ +public interface Ai1MemoryService { + + /** + * 修改长期记忆正文 + * + * @param updateReqVO 修改信息 + */ + void updateMemory(@Valid Ai1MemoryUpdateReqVO updateReqVO); + + /** + * 删除长期记忆（逻辑删除）；之后的归档会根据历史消息生成新的记录 + * + * @param id 编号 + */ + void deleteMemory(Long id); + + /** + * 批量删除长期记忆 + * + * @param ids 编号列表 + */ + void deleteMemoryListByIds(List ids); + + /** + * 删除 Agent 下的全部长期记忆（Agent 删除时调用） + * + * @param agentIds Agent 编号集合 + */ + void deleteMemoryListByAgentIds(Collection agentIds); + + /** + * 获得长期记忆 + * + * @param id 编号 + * @return 长期记忆 + */ + Ai1MemoryDO getMemory(Long id); + + /** + * 获得用户在 Agent 下的长期记忆 + * + * @param agentId Agent 编号 + * @param userId 用户编号 + * @return 长期记忆 + */ + Ai1MemoryDO getMemoryByAgentIdAndUserId(Long agentId, Long userId); + + /** + * 创建归档的长期记忆；写入前校验，Agent 已删除，或同一用户已有未删除的记忆时不创建 + * + * @param memory 长期记忆 + * @return 是否创建成功 + */ + boolean createMemoryIfAgentExists(Ai1MemoryDO memory); + + /** + * 更新归档的长期记忆；写入前校验，读取后被人工修改或删除时不更新 + * + * @param old 归档读取时的记忆 + * @param content 新的记忆正文 + * @param lastMessageId 新的水位线 + * @return 是否更新成功 + */ + boolean updateMemoryIfUnchanged(Ai1MemoryDO old, String content, Long lastMessageId); + + /** + * 获得长期记忆分页 + * + * @param pageReqVO 分页查询 + * @return 长期记忆分页 + */ + PageResult getMemoryPage(Ai1MemoryPageReqVO pageReqVO); + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java new file mode 100644 index 0000000..8da768e --- /dev/null +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java @@ -0,0 +1,123 @@ +package cn.iocoder.yudao.module.ai1.service.memory; + +import cn.hutool.core.collection.CollUtil; +import cn.hutool.core.util.ObjUtil; +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.framework.common.util.object.BeanUtils; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.mysql.memory.Ai1MemoryMapper; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; +import jakarta.annotation.Resource; +import org.springframework.context.annotation.Lazy; +import org.springframework.stereotype.Service; +import org.springframework.validation.annotation.Validated; + +import java.util.Collection; +import java.util.Collections; +import java.util.List; + +import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception; +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MEMORY_NOT_EXISTS; + +/** + * AI1 长期记忆 Service 实现类 + * + * @author 芋道源码 + */ +@Service +@Validated +public class Ai1MemoryServiceImpl implements Ai1MemoryService { + + @Resource + private Ai1MemoryMapper memoryMapper; + + @Resource + @Lazy // 延迟加载，避免循环依赖 + private Ai1AgentService agentService; + + @Override + public void updateMemory(Ai1MemoryUpdateReqVO updateReqVO) { + // 1. 校验存在 + validateMemoryExists(updateReqVO.getId()); + + // 2. 更新正文；水位线不变，之后的归档写入前校验到正文变化时放弃 + Ai1MemoryDO updateObj = BeanUtils.toBean(updateReqVO, Ai1MemoryDO.class); + memoryMapper.updateById(updateObj); + } + + @Override + public void deleteMemory(Long id) { + deleteMemoryListByIds(Collections.singletonList(id)); + } + + @Override + public void deleteMemoryListByIds(List ids) { + // 1. 校验存在 + if (CollUtil.isEmpty(ids)) { + return; + } + ids.forEach(this::validateMemoryExists); + + // 2. 逻辑删除，保留历史记录；进行中的归档写入会因记录已删除而放弃 + memoryMapper.deleteByIds(ids); + } + + @Override + public void deleteMemoryListByAgentIds(Collection agentIds) { + if (CollUtil.isEmpty(agentIds)) { + return; + } + memoryMapper.deleteBatch(Ai1MemoryDO::getAgentId, agentIds); + } + + private void validateMemoryExists(Long id) { + if (memoryMapper.selectById(id) == null) { + throw exception(MEMORY_NOT_EXISTS); + } + } + + @Override + public Ai1MemoryDO getMemory(Long id) { + return memoryMapper.selectById(id); + } + + @Override + public Ai1MemoryDO getMemoryByAgentIdAndUserId(Long agentId, Long userId) { + return memoryMapper.selectByAgentIdAndUserId(agentId, userId); + } + + @Override + public boolean createMemoryIfAgentExists(Ai1MemoryDO memory) { + // 1. 前置校验 Agent 存在，且该用户没有未删除的记忆；同租户归档由归档锁串行，极端并发下的重复记录由按用户查询容忍 + if (agentService.getAgent(memory.getAgentId()) == null + || memoryMapper.selectByAgentIdAndUserId(memory.getAgentId(), memory.getUserId()) != null) { + return false; + } + + // 2. 插入 + memoryMapper.insert(memory); + return true; + } + + @Override + public boolean updateMemoryIfUnchanged(Ai1MemoryDO old, String content, Long lastMessageId) { + // 1. 校验记忆未删除，且水位线与正文都未变化 + Ai1MemoryDO memory = memoryMapper.selectById(old.getId()); + if (memory == null || ObjUtil.notEqual(memory.getLastMessageId(), old.getLastMessageId()) + || ObjUtil.notEqual(memory.getContent(), old.getContent())) { + return false; + } + + // 2. 更新正文并推进水位线；校验之后的人工修改不再拦截，已删除的记录不会被更新 + return memoryMapper.updateById(new Ai1MemoryDO().setId(old.getId()) + .setContent(content).setLastMessageId(lastMessageId)) > 0; + } + + @Override + public PageResult getMemoryPage(Ai1MemoryPageReqVO pageReqVO) { + return memoryMapper.selectPage(pageReqVO); + } + +} diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java index 61c3170..66bc104 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java @@ -1,5 +1,7 @@ package cn.iocoder.yudao.module.ai1.service.session; +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import jakarta.validation.Valid; @@ -24,6 +26,14 @@ public interface Ai1SessionMessageService { List getSessionMessageListBySessionId(Long sessionId); /** + * 获得会话消息分页，按编号倒序 + * + * @param pageReqVO 分页查询 + * @return 消息分页 + */ + PageResult getSessionMessagePage(Ai1SessionMessagePageReqVO pageReqVO); + + /** * 获得会话中指定消息之前最近的若干条消息，按编号倒序 * * @param sessionId 会话编号 @@ -34,6 +44,37 @@ public interface Ai1SessionMessageService { List getSessionMessageListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit); /** + * 获得会话中生成中的助手消息编号列表 + * + * @param sessionId 会话编号 + * @return 消息编号列表 + */ + List getGeneratingAssistantMessageIdList(Long sessionId); + + /** + * 获得多个会话中，编号大于 id 的第一条指定状态的消息 + * + * @param sessionIds 会话编号集合 + * @param status 状态 + * @param id 消息编号（不含） + * @return 消息 + */ + Ai1SessionMessageDO getFirstSessionMessageBySessionIdsAndStatusAndIdGt(Collection sessionIds, Integer status, Long id); + + /** + * 获得多个会话中，编号在 (minId, maxId) 之间的指定状态消息，按编号升序 + * + * @param sessionIds 会话编号集合 + * @param status 状态 + * @param minId 最小编号（不含） + * @param maxId 最大编号（不含），为空时不限 + * @param limit 数量上限 + * @return 消息列表 + */ + List getSessionMessageListBySessionIdsAndStatusAndIdBetween(Collection sessionIds, Integer status, + Long minId, Long maxId, Integer limit); + + /** * 发送消息（SSE 流式）：校验会话归属 → 落库用户消息与助手占位 → 投递生成任务 → 打开 SSE 连接 * * 校验、落库失败时不抛出异常，而是返回只包含 error 事件的 SSE 连接 @@ -55,6 +96,14 @@ public interface Ai1SessionMessageService { SseEmitter resumeSessionMessageStream(Long userId, Long messageId, String lastEventId); /** + * 停止【我的】会话生成：校验会话归属后，终止会话中生成中的助手消息，跨节点生效 + * + * @param userId 用户编号 + * @param sessionId 会话编号 + */ + void stopSessionMessageMy(Long userId, Long sessionId); + + /** * 获得消息 * * @param id 编号 diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImpl.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImpl.java index d46e724..aa37ec6 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImpl.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImpl.java @@ -4,11 +4,14 @@ import cn.hutool.core.collection.CollUtil; import cn.hutool.core.util.StrUtil; import cn.hutool.extra.spring.SpringUtil; import cn.iocoder.yudao.framework.common.exception.ServiceException; +import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMessageMapper; +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; @@ -21,6 +24,7 @@ import org.springframework.validation.annotation.Validated; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; import java.util.Collection; +import java.util.Collections; import java.util.List; import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception; @@ -49,25 +53,57 @@ public class Ai1SessionMessageServiceImpl implements Ai1SessionMessageService { @Resource private Ai1SessionStreamTool sessionStreamTool; + @Resource + private Ai1MemoryLockRedisDAO memoryLockRedisDAO; + @Override public List getSessionMessageListBySessionId(Long sessionId) { return sessionMessageMapper.selectListBySessionId(sessionId); } @Override + public PageResult getSessionMessagePage(Ai1SessionMessagePageReqVO pageReqVO) { + return sessionMessageMapper.selectPage(pageReqVO); + } + + @Override public List getSessionMessageListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit) { return sessionMessageMapper.selectListBySessionIdAndIdLessThan(sessionId, id, limit); } @Override + public List getGeneratingAssistantMessageIdList(Long sessionId) { + return sessionMessageMapper.selectIdListBySessionIdAndRoleAndStatus(sessionId, + Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + } + + @Override + public Ai1SessionMessageDO getFirstSessionMessageBySessionIdsAndStatusAndIdGt(Collection sessionIds, Integer status, Long id) { + if (CollUtil.isEmpty(sessionIds)) { + return null; + } + return sessionMessageMapper.selectFirstBySessionIdsAndStatusAndIdGt(sessionIds, status, id); + } + + @Override + public List getSessionMessageListBySessionIdsAndStatusAndIdBetween(Collection sessionIds, Integer status, + Long minId, Long maxId, Integer limit) { + if (CollUtil.isEmpty(sessionIds)) { + return Collections.emptyList(); + } + return sessionMessageMapper.selectListBySessionIdsAndStatusAndIdBetween(sessionIds, status, minId, maxId, limit); + } + + @Override public SseEmitter sendSessionMessageStream(Long userId, Ai1SessionMessageSendReqVO sendReqVO) { try { // 1. 校验会话归属、Agent 存在且已开启 Ai1SessionDO session = sessionService.validateSessionMy(userId, sendReqVO.getSessionId()); agentService.validateAgentEnabled(session.getAgentId()); - // 2. 同一事务内落库用户消息与助手占位 - Long messageId = getSelf().createRoundMessages(session, sendReqVO.getContent()); + // 2. 同一事务内落库用户消息与助手占位；持有记忆快照锁至事务提交 + Long messageId = memoryLockRedisDAO.lockSnapshot(session.getAgentId(), session.getUserId(), + () -> getSelf().createRoundMessages(session, sendReqVO.getContent())); // 3. 投递生成任务，并打开 SSE 连接 try { @@ -120,6 +156,14 @@ public class Ai1SessionMessageServiceImpl implements Ai1SessionMessageService { } } + @Override + public void stopSessionMessageMy(Long userId, Long sessionId) { + // 1. 校验会话属于当前用户 + sessionService.validateSessionMy(userId, sessionId); + // 2. 取消会话中生成中的助手消息 + sessionStreamTool.cancel(sessionId); + } + /** * 开启一轮会话：首条消息自动生成标题 → 落库用户消息（完成）与助手占位（生成中）→ 刷新会话活跃时间 * diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java index e034443..240e605 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java @@ -54,6 +54,14 @@ public interface Ai1SessionService { Ai1SessionDO validateSessionMy(Long userId, Long id); /** + * 获得会话 + * + * @param id 编号 + * @return 会话 + */ + Ai1SessionDO getSession(Long id); + + /** * 获得指定用户在指定 Agent 下的会话列表，按编号倒序 * * @param agentId Agent 编号 @@ -63,6 +71,14 @@ public interface Ai1SessionService { List getSessionListByAgentIdAndUserId(Long agentId, Long userId); /** + * 获得 Agent 下有会话的用户编号列表 + * + * @param agentId Agent 编号 + * @return 用户编号列表 + */ + List getUserIdListByAgentId(Long agentId); + + /** * 获得会话分页（管理端查看 Agent 的全部会话） * * @param pageReqVO 分页查询 diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImpl.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImpl.java index 96d31d6..723eaa2 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImpl.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImpl.java @@ -84,11 +84,21 @@ public class Ai1SessionServiceImpl implements Ai1SessionService { } @Override + public Ai1SessionDO getSession(Long id) { + return sessionMapper.selectById(id); + } + + @Override public List getSessionListByAgentIdAndUserId(Long agentId, Long userId) { return sessionMapper.selectListByAgentIdAndUserId(agentId, userId); } @Override + public List getUserIdListByAgentId(Long agentId) { + return sessionMapper.selectUserIdListByAgentId(agentId); + } + + @Override public PageResult getSessionPage(Ai1SessionPageReqVO pageReqVO) { return sessionMapper.selectPage(pageReqVO); } diff --git a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java index a9d239e..aa63746 100644 --- a/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java +++ b/yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java @@ -87,4 +87,56 @@ public class Ai1Utils { } } + /** + * 规范化模型生成的长期记忆 Markdown：修正标题、列表标记，并保证维度标题与内容不在同一行 + * + * 只按长期记忆的固定维度标题拆行，不拆开 C#、#123、URL 锚点，不把中文破折号视为列表标记 + * + * @param content 模型输出的记忆 + * @param sections 记忆维度标题 + * @return 规范化后的记忆；为空时原样返回 + */ + public static String normalizeMemoryMarkdown(String content, Collection sections) { + if (StrUtil.isBlank(content)) { + return content; + } + String text = content.replace("\r\n", "\n").replace("\r", "\n").trim(); + // 1. 行内出现的维度标题另起一行：…文字## 用户画像 → …文字\n## 用户画像 + text = text.replaceAll("([^#\\n])(#{1,6}[ \\t]*(?:" + String.join("|", sections) + "))", "$1\n$2"); + // 2. 标题标记后补空格：#标题 → # 标题（# 后为非 #、非空白，避免把 ## 拆成 # #） + text = text.replaceAll("(?m)^(#{1,6})(?=[^#\\s])", "$1 "); + // 3. 维度标题与紧随其后的内容拆行：## 偏好习惯-xxx → ## 偏好习惯\n-xxx + for (String section : sections) { + text = text.replaceAll("(?m)(#{1,6} " + section + ")[ \\t]*(?=\\S)", "$1\n"); + } + // 4. 行内列表项另起一行，仅当前一字符为非 ASCII：…口吻-开场 → …口吻\n- 开场（不把中文破折号 —、– 视为列表标记） + text = text.replaceAll("([^\\x00-\\x7F])[-*+\\u2022][ \\t]*(?=\\S)", "$1\n- "); + // 5. 行首列表标记后补空格：-xxx → - xxx（同样不处理中文破折号） + text = text.replaceAll("(?m)^([-*+\\u2022])(?=\\S)", "$1 "); + // 6. 不同维度之间空一行 + text = text.replaceAll("(?m)^(#{1,6} )", "\n$1"); + // 7. 合并多余空行 + text = text.replaceAll("\\n{3,}", "\n\n"); + return text.trim(); + } + + /** + * 去除首尾空白后按字符数截断，不拆开 UTF-16 代理对 + * + * @param content 内容 + * @param maxLength 最大长度（字符），需大于 0 + * @return 截断后的内容；为空时原样返回 + */ + public static String limitLength(String content, int maxLength) { + if (StrUtil.isBlank(content)) { + return content; + } + String text = content.trim(); + if (text.length() <= maxLength) { + return text; + } + int end = Character.isHighSurrogate(text.charAt(maxLength - 1)) ? maxLength - 1 : maxLength; + return text.substring(0, end); + } + } diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatToolTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatToolTest.java index 67d51f2..15b58bc 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatToolTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatToolTest.java @@ -18,7 +18,10 @@ import org.springframework.ai.tool.annotation.Tool; import org.springframework.ai.tool.function.FunctionToolCallback; import reactor.core.publisher.Flux; +import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.*; @@ -89,6 +92,41 @@ public class Ai1LlmChatToolTest extends BaseMockitoUnitTest { assertEquals(1, skillCalls.get()); } + @Test + public void testChat_whitespaceChunks() { + // mock chatModel 和 llmModelFactory 的方法：空格、换行分片也是有效增量，只跳过空串 + ChatModel chatModel = mock(ChatModel.class); + when(chatModel.getDefaultOptions()).thenReturn(ToolCallingChatOptions.builder().build()); + when(chatModel.stream(any(Prompt.class))).thenReturn(Flux.just( + buildChunk("想", "你好"), buildChunk(" ", " "), buildChunk("\n", "\n"), + buildChunk("", "世界"), buildChunk(null, ""))); + ChatClient chatClient = ChatClient.create(chatModel); + when(llmModelFactory.buildChatClient(any(), any())).thenReturn(chatClient); + // 准备参数 + List thinkingDeltas = new ArrayList<>(); + List contentDeltas = new ArrayList<>(); + + // 调用 + Ai1LlmChatTool.ChatText result = llmChatTool.chat(new Ai1ModelRespBO(), "", List.of(), "你好", + List.of(), List.of(), 1L, thinkingDeltas::add, contentDeltas::add); + // 断言 + assertEquals("你好 \n世界", result.getContent()); + assertEquals("想 \n", result.getThinking()); + assertEquals(List.of("你好", " ", "\n", "世界"), contentDeltas); + assertEquals(List.of("想", " ", "\n"), thinkingDeltas); + } + + // ========== 测试数据 ========== + + private static ChatResponse buildChunk(String reasoning, String content) { + Map properties = new HashMap<>(); + if (reasoning != null) { + properties.put("reasoningContent", reasoning); + } + return new ChatResponse(List.of(new Generation(AssistantMessage.builder() + .content(content).properties(properties).build()))); + } + // ========== 测试工具 ========== public static class AnnotatedTools { diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java new file mode 100644 index 0000000..3453272 --- /dev/null +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java @@ -0,0 +1,294 @@ +package cn.iocoder.yudao.module.ai1.harness.memory; + +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; +import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder; +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO; +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; +import cn.iocoder.yudao.module.ai1.harness.session.Ai1SessionStreamTool; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService; +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.InOrder; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.Spy; +import org.redisson.api.RLock; +import org.redisson.api.RedissonClient; +import org.springframework.test.util.ReflectionTestUtils; + +import java.util.List; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +/** + * {@link Ai1MemoryArchiveTool} 的单元测试 + * + * @author 芋道源码 + */ +public class Ai1MemoryArchiveToolTest extends BaseMockitoUnitTest { + + private static final Long TENANT_ID = 1L; + private static final Long AGENT_ID = 1L; + private static final Long USER_ID = 100L; + private static final List SESSION_IDS = List.of(10L, 11L); + + @InjectMocks + private Ai1MemoryArchiveTool memoryArchiveTool; + + @Spy + private Ai1MemoryLockRedisDAO memoryLockRedisDAO; + @Mock + private RedissonClient redissonClient; + @Mock + private RLock archiveLock; + @Mock + private RLock snapshotLock; + @Spy + private YudaoAi1Properties ai1Properties = new YudaoAi1Properties(); + + @Mock + private Ai1AgentService agentService; + @Mock + private Ai1SessionService sessionService; + @Mock + private Ai1SessionMessageService sessionMessageService; + @Mock + private Ai1MemoryService memoryService; + @Mock + private Ai1MemoryTool memoryTool; + @Mock + private Ai1SessionStreamTool sessionStreamTool; + + @BeforeEach + public void before() { + TenantContextHolder.setTenantId(TENANT_ID); + // mock 归档锁、快照锁 + ReflectionTestUtils.setField(memoryLockRedisDAO, "redissonClient", redissonClient); + lenient().when(redissonClient.getLock("ai1:memory:archive-lock:" + TENANT_ID)).thenReturn(archiveLock); + lenient().when(archiveLock.tryLock()).thenReturn(true); + lenient().when(archiveLock.isHeldByCurrentThread()).thenReturn(true); + lenient().when(redissonClient.getLock("ai1:memory:snapshot-lock:" + AGENT_ID + ":" + USER_ID)).thenReturn(snapshotLock); + lenient().when(sessionService.getSessionListByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(List.of( + new Ai1SessionDO().setId(10L), new Ai1SessionDO().setId(11L))); + } + + @AfterEach + public void after() { + TenantContextHolder.clear(); + } + + @Test + public void testArchive_locked() { + // mock 其他节点归档中 + when(archiveLock.tryLock()).thenReturn(false); + + // 调用，并断言 + assertEquals(0, memoryArchiveTool.archive()); + verifyNoInteractions(agentService, memoryTool); + verify(archiveLock, never()).unlock(); + } + + @Test + public void testArchive_lockNotHeld() { + // mock 方法：归档结束时锁已不属于当前线程 + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of()); + when(archiveLock.isHeldByCurrentThread()).thenReturn(false); + + // 调用，并断言：不释放其他线程持有的锁 + assertEquals(0, memoryArchiveTool.archive()); + verify(archiveLock, never()).unlock(); + } + + @Test + public void testArchive_onlyMemoryEnabled() { + // mock 方法：只有开启记忆的 Agent 归档 + Ai1AgentDO agent = buildAgent(); + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of( + agent, buildAgent().setId(2L).setMemoryEnabled(false), buildAgent().setId(3L).setMemoryEnabled(null))); + when(sessionService.getUserIdListByAgentId(AGENT_ID)).thenReturn(List.of(USER_ID)); + mockNewMessages(0L, null, List.of(buildMessage(12L), buildMessage(13L))); + when(memoryTool.generate(eq(agent), isNull(), anyList(), eq(2000))).thenReturn("## 用户画像\n- 程序员"); + when(memoryService.createMemoryIfAgentExists(any())).thenReturn(true); + + // 调用 + int count = memoryArchiveTool.archive(); + // 断言：首次归档插入，水位线为本批最后一条 + assertEquals(1, count); + verify(sessionService, never()).getUserIdListByAgentId(2L); + verify(sessionService, never()).getUserIdListByAgentId(3L); + verify(memoryService).createMemoryIfAgentExists(argThat(memory -> AGENT_ID.equals(memory.getAgentId()) + && USER_ID.equals(memory.getUserId()) && "## 用户画像\n- 程序员".equals(memory.getContent()) + && Long.valueOf(13L).equals(memory.getLastMessageId()))); + verify(archiveLock).unlock(); + } + + @Test + public void testArchive_userFailedNotAffectOthers() { + // mock 方法：第一个用户生成失败，第二个用户无新消息 + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of(buildAgent())); + when(sessionService.getUserIdListByAgentId(AGENT_ID)).thenReturn(List.of(USER_ID, 200L)); + mockNewMessages(0L, null, List.of(buildMessage(12L))); + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenThrow(new IllegalStateException("模型超时")); + when(redissonClient.getLock("ai1:memory:snapshot-lock:" + AGENT_ID + ":200")).thenReturn(snapshotLock); + + // 调用 + int count = memoryArchiveTool.archive(); + // 断言：不推进水位线，下次重试 + assertEquals(0, count); + verify(memoryService).getMemoryByAgentIdAndUserId(AGENT_ID, 200L); + verify(memoryService, never()).createMemoryIfAgentExists(any()); + verify(memoryService, never()).updateMemoryIfUnchanged(any(), any(), any()); + verify(archiveLock).unlock(); + } + + @Test + public void testArchiveOne_safePrefix() { + // mock 方法：水位线之后有生成中的消息，只读取它之前的已完成消息 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).agentId(AGENT_ID).userId(USER_ID) + .content("旧记忆").lastMessageId(11L).build(); + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory); + mockNewMessages(11L, 20L, List.of(buildMessage(12L), buildMessage(13L))); + when(memoryTool.generate(any(), eq("旧记忆"), anyList(), eq(2000))).thenReturn("新记忆"); + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 13L)).thenReturn(true); + + // 调用，并断言 + assertTrue(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + // 断言：孤儿收尾在快照锁外，持有快照锁读取，模型在锁外调用 + InOrder inOrder = inOrder(sessionStreamTool, snapshotLock, sessionMessageService, memoryTool); + inOrder.verify(sessionStreamTool).failStaleMessage(argThat(message -> Long.valueOf(20L).equals(message.getId()))); + inOrder.verify(snapshotLock).lock(); + inOrder.verify(sessionMessageService).getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 11L, 20L, Ai1MemoryArchiveTool.BATCH_SIZE); + inOrder.verify(snapshotLock).unlock(); + inOrder.verify(memoryTool).generate(any(), any(), anyList(), anyInt()); + } + + @Test + public void testArchiveOne_snapshotReadFailed() { + // mock 方法：持有快照锁读取新消息时失败 + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 0L, null, Ai1MemoryArchiveTool.BATCH_SIZE)) + .thenThrow(new IllegalStateException("数据库超时")); + + // 调用，并断言：异常抛出，快照锁仍释放 + assertThrows(IllegalStateException.class, () -> memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verify(snapshotLock).lock(); + verify(snapshotLock).unlock(); + verifyNoInteractions(memoryTool); + } + + @Test + public void testArchiveOne_noNewMessages() { + // mock 方法 + mockNewMessages(0L, null, List.of()); + + // 调用，并断言 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verifyNoInteractions(memoryTool, sessionStreamTool); + } + + @Test + public void testArchiveOne_staleGeneratingRecovered() { + // mock 数据：孤儿生成之后有已完成的消息 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).content("旧记忆").lastMessageId(11L).build(); + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory); + // mock 方法：孤儿生成收尾后，不再有生成中的消息 + Ai1SessionMessageDO generatingMessage = buildMessage(20L); + when(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(SESSION_IDS, + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), 11L)).thenReturn(generatingMessage, (Ai1SessionMessageDO) null); + when(sessionStreamTool.failStaleMessage(generatingMessage)).thenReturn(true); + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 11L, null, Ai1MemoryArchiveTool.BATCH_SIZE)) + .thenReturn(List.of(buildMessage(12L), buildMessage(21L))); + when(memoryTool.generate(any(), eq("旧记忆"), anyList(), eq(2000))).thenReturn("新记忆"); + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 21L)).thenReturn(true); + + // 调用，并断言：水位线越过已收尾的孤儿生成 + assertTrue(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + InOrder inOrder = inOrder(sessionStreamTool, snapshotLock); + inOrder.verify(sessionStreamTool).failStaleMessage(generatingMessage); + inOrder.verify(snapshotLock).lock(); + } + + @Test + public void testArchiveOne_activeGeneratingBlocks() { + // mock 方法：生成仍在执行或未超时，安全前缀为空 + mockNewMessages(0L, 12L, List.of()); + + // 调用，并断言：不收尾，也不调用模型 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verify(sessionStreamTool).failStaleMessage(argThat(message -> Long.valueOf(12L).equals(message.getId()))); + verifyNoInteractions(memoryTool); + } + + @Test + public void testArchiveOne_blankMemory() { + // mock 方法 + mockNewMessages(0L, null, List.of(buildMessage(12L))); + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn(" "); + + // 调用，并断言：模型无输出时不写入，水位线不前进 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verify(memoryService, never()).createMemoryIfAgentExists(any()); + } + + @Test + public void testArchiveOne_changedDuringArchive() { + // mock 方法：归档期间记忆被人工修改或删除 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).content("旧记忆").lastMessageId(11L).build(); + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory); + mockNewMessages(11L, null, List.of(buildMessage(12L))); + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn("新记忆"); + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 12L)).thenReturn(false); + + // 调用，并断言：放弃本次结果，不重新插入 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verify(memoryService, never()).createMemoryIfAgentExists(any()); + } + + @Test + public void testArchiveOne_insertConflict() { + // mock 方法：其他归档已插入，或 Agent 已删除 + mockNewMessages(0L, null, List.of(buildMessage(12L))); + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn("新记忆"); + when(memoryService.createMemoryIfAgentExists(any())).thenReturn(false); + + // 调用，并断言：放弃本次结果 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000)); + verify(memoryService, never()).updateMemoryIfUnchanged(any(), any(), any()); + } + + // ========== mock 方法 ========== + + private void mockNewMessages(Long lastMessageId, Long generatingMessageId, List messages) { + when(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(SESSION_IDS, + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId)) + .thenReturn(generatingMessageId != null ? buildMessage(generatingMessageId) : null); + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), lastMessageId, generatingMessageId, Ai1MemoryArchiveTool.BATCH_SIZE)) + .thenReturn(messages); + } + + // ========== 测试数据 ========== + + private static Ai1AgentDO buildAgent() { + return new Ai1AgentDO().setId(AGENT_ID).setProviderId(1L).setModelId(2L).setMemoryEnabled(true); + } + + private static Ai1SessionMessageDO buildMessage(Long id) { + return new Ai1SessionMessageDO().setId(id).setContent("消息" + id); + } + +} diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java new file mode 100644 index 0000000..7e8e5bf --- /dev/null +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java @@ -0,0 +1,149 @@ +package cn.iocoder.yudao.module.ai1.harness.memory; + +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; +import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; +import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; +import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService; +import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; +import org.mockito.InjectMocks; +import org.mockito.Mock; + +import java.util.List; + +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException; +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MODEL_TYPE_NOT_CHAT; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +/** + * {@link Ai1MemoryTool} 的单元测试 + * + * @author 芋道源码 + */ +public class Ai1MemoryToolTest extends BaseMockitoUnitTest { + + @InjectMocks + private Ai1MemoryTool memoryTool; + + @Mock + private Ai1MemoryService memoryService; + @Mock + private Ai1ModelService modelService; + @Mock + private Ai1LlmChatTool llmChatTool; + + @Test + public void testLoadMemory() { + // mock 方法 + when(memoryService.getMemoryByAgentIdAndUserId(1L, 100L)).thenReturn(Ai1MemoryDO.builder().content("## 用户画像\n- 程序员").build()); + + // 调用，并断言 + assertEquals("## 用户画像\n- 程序员", memoryTool.loadMemory(1L, 100L)); + // 调用，并断言：编号为空时不查询 + assertNull(memoryTool.loadMemory(null, 100L)); + verify(memoryService, never()).getMemoryByAgentIdAndUserId(isNull(), any()); + } + + @Test + public void testMergePrompt() { + // 调用，并断言：按模板合入记忆 + assertEquals("你是客服\n\n以下是该用户的历史长期记忆（Markdown），请在回答时结合这些信息；" + + "若与用户当前表述冲突，以当前表述为准：\n## 用户画像\n- 程序员", memoryTool.mergePrompt("你是客服", " ## 用户画像\n- 程序员 \n")); + // 调用，并断言：记忆为空时原样返回 + assertEquals("你是客服", memoryTool.mergePrompt("你是客服", " ")); + assertEquals("", memoryTool.mergePrompt(null, null)); + } + + @Test + public void testGenerate_modelNotChat() { + // mock modelService 的方法 + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO() + .setModelType(Ai1ModelTypeEnum.EMBEDDING.getType())); + + // 调用，并断言异常 + assertServiceException(() -> memoryTool.generate(buildAgent(), null, List.of(), 2000), MODEL_TYPE_NOT_CHAT); + verifyNoInteractions(llmChatTool); + } + + @Test + public void testGenerate_prompt() { + // mock 方法 + mockChat("## 用户画像\n- 程序员"); + // 准备参数 + List messages = List.of( + new Ai1SessionMessageDO().setRole(Ai1SessionMessageRoleEnum.USER.getRole()).setContent("我是程序员"), + new Ai1SessionMessageDO().setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()).setContent("好的")); + + // 调用 + String memory = memoryTool.generate(buildAgent(), " 旧记忆 ", messages, 500); + // 断言：单轮会话，不带历史、工具、RAG，也不替换会话 Header + assertEquals("## 用户画像\n- 程序员", memory); + ArgumentCaptor systemPromptCaptor = ArgumentCaptor.forClass(String.class); + ArgumentCaptor userPromptCaptor = ArgumentCaptor.forClass(String.class); + verify(llmChatTool).chat(any(), systemPromptCaptor.capture(), eq(List.of()), userPromptCaptor.capture(), + eq(List.of()), eq(List.of()), isNull(), isNull(), isNull()); + assertTrue(systemPromptCaptor.getValue().endsWith("总长度不超过 500 字。")); + assertEquals("【历史记忆】\n旧记忆\n\n【最新对话】\n用户：我是程序员\n助手：好的\n", userPromptCaptor.getValue()); + } + + @Test + public void testGenerate_normalizeSections() { + // mock 方法：标题与内容压成一行 + mockChat("##用户画像-后端工程师##偏好习惯-简洁回答"); + + // 调用 + String memory = memoryTool.generate(buildAgent(), null, List.of(), 2000); + // 断言 + assertEquals("## 用户画像\n- 后端工程师\n\n## 偏好习惯\n- 简洁回答", memory); + } + + @Test + public void testGenerate_defaultMaxLength() { + // mock 方法 + mockChat("用户喜欢😀表情"); + + // 调用：未配置长度上限时使用默认值 + String memory = memoryTool.generate(buildAgent(), null, List.of(), 0); + // 断言 + assertEquals("用户喜欢😀表情", memory); + ArgumentCaptor systemPromptCaptor = ArgumentCaptor.forClass(String.class); + verify(llmChatTool).chat(any(), systemPromptCaptor.capture(), anyList(), anyString(), anyList(), anyList(), + isNull(), isNull(), isNull()); + assertTrue(systemPromptCaptor.getValue().endsWith("总长度不超过 2000 字。")); + // 调用，并断言：按长度上限截断 + assertEquals("用户喜欢", memoryTool.generate(buildAgent(), null, List.of(), 5)); + } + + @Test + public void testGenerate_empty() { + // mock 方法 + mockChat(" "); + + // 调用，并断言：模型无输出时返回空，由调用方跳过写入 + assertTrue(memoryTool.generate(buildAgent(), null, List.of(), 2000).isBlank()); + } + + // ========== mock 方法 ========== + + private void mockChat(String content) { + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO() + .setModelType(Ai1ModelTypeEnum.CHAT.getType())); + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), isNull(), isNull(), isNull())) + .thenReturn(new Ai1LlmChatTool.ChatText(content, "")); + } + + // ========== 测试数据 ========== + + private static Ai1AgentDO buildAgent() { + return new Ai1AgentDO().setId(1L).setProviderId(1L).setModelId(2L); + } + +} diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamToolTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamToolTest.java index 2f57288..ca4adda 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamToolTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionStreamToolTest.java @@ -3,6 +3,7 @@ package cn.iocoder.yudao.module.ai1.harness.session; import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder; import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest; import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; @@ -10,7 +11,9 @@ import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool; import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory; +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool; import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool; +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker; import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory; import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; import cn.iocoder.yudao.module.ai1.service.knowledge.Ai1KnowledgeBaseService; @@ -20,29 +23,52 @@ import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.OS; +import org.junit.jupiter.api.io.TempDir; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import org.mockito.ArgumentCaptor; import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.Spy; +import org.redisson.api.RBucket; import org.redisson.api.RLock; import org.redisson.api.RStream; +import org.redisson.api.RTopic; import org.redisson.api.RedissonClient; import org.redisson.api.stream.StreamAddArgs; +import org.redisson.api.stream.StreamCreateGroupArgs; import org.redisson.api.stream.StreamMessageId; import org.redisson.api.stream.StreamRangeArgs; import org.redisson.api.stream.StreamReadArgs; +import org.redisson.api.stream.StreamReadGroupArgs; +import org.redisson.client.RedisException; import org.redisson.client.codec.Codec; +import org.springaicommunity.agent.common.exec.ExecBackend; +import org.springaicommunity.agent.common.exec.ExecHandle; +import org.springaicommunity.agent.common.exec.ExecSpec; +import org.springaicommunity.agent.exec.LocalExecBackend; import org.springframework.test.util.ReflectionTestUtils; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.time.Duration; +import java.time.LocalDateTime; import java.util.ArrayList; import java.util.HashMap; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CancellationException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.function.Consumer; @@ -74,6 +100,10 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { private RStream resultStream; @Mock private RLock taskLock; + @Mock + private RBucket stopBucket; + @Mock + private RTopic stopTopic; @Spy private YudaoAi1Properties ai1Properties = new YudaoAi1Properties(); @@ -95,15 +125,24 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { private Ai1SkillToolFactory skillToolFactory; @Mock private Ai1RagTool ragTool; + @Mock + private Ai1MemoryTool memoryTool; + @Mock + private Ai1SkillShellTracker shellTracker; @BeforeEach public void before() { + // 直接调用 handleTask 的用例按运行中处理，与 start 之后一致 + ReflectionTestUtils.setField(sessionStreamTool, "running", true); // mock redissonClient 的方法（部分用例不访问 Redis） lenient().doReturn(taskStream).when(redissonClient).getStream(eq("ai1:session:tasks"), any(Codec.class)); lenient().doReturn(resultStream).when(redissonClient).getStream(eq("ai1:session:result:" + MESSAGE_ID), any(Codec.class)); lenient().when(redissonClient.getLock("ai1:session:task-lock:" + MESSAGE_ID)).thenReturn(taskLock); lenient().when(taskLock.tryLock()).thenReturn(true); lenient().when(taskLock.isHeldByThread(anyLong())).thenReturn(true); + lenient().doReturn(stopBucket).when(redissonClient).getBucket(eq("ai1:session:stop:" + MESSAGE_ID), any(Codec.class)); + lenient().when(redissonClient.getTopic(eq("ai1:session:stop"), any(Codec.class))).thenReturn(stopTopic); + lenient().when(skillToolFactory.newShellTracker(AGENT_ID)).thenReturn(shellTracker); lenient().when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO() .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus())); @@ -256,11 +295,14 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { // 调用 sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); - // 断言 + // 断言：正文追加固定失败提示，不拼接原始异常 ArgumentCaptor messageCaptor = ArgumentCaptor.forClass(Ai1SessionMessageDO.class); verify(sessionMessageService).updateSessionMessage(messageCaptor.capture()); assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), messageCaptor.getValue().getStatus()); + assertEquals("\n\n> 本条回复生成失败", messageCaptor.getValue().getContent()); assertEquals(List.of("ping", "error", "done"), captureResultTypes()); + assertEquals(List.of("", "本条回复生成失败", ""), captureResultData()); + verify(shellTracker).cancel(); } @Test @@ -310,7 +352,7 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), reclaimed); // 断言 verifyNoInteractions(llmChatTool, agentService); - verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段第二段".equals(message.getContent()) + verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段第二段\n\n> 生成任务已中断，请重新发送".equals(message.getContent()) && "思考".equals(message.getReasoning()) && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); assertEquals(List.of("error", "done"), captureResultTypes()); @@ -326,7 +368,7 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true); // 断言 verifyNoInteractions(llmChatTool); - verify(sessionMessageService).updateSessionMessage(argThat(message -> "".equals(message.getContent()) + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 生成任务已中断，请重新发送".equals(message.getContent()) && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); assertEquals(List.of("error", "done"), captureResultTypes()); } @@ -437,6 +479,573 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { verify(taskStream).ack("ai1-session-workers", RECORD_ID); } + @Test + public void testHandleTask_terminatedWithoutTerminal() { + // mock 方法：人工终止已回填，done 尚未写入 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO() + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent("已生成\n\n> 本条回复已被人工终止")); + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of()); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true); + // 断言：人工终止不补写 error + verifyNoInteractions(llmChatTool); + verify(sessionMessageService, never()).updateSessionMessage(any()); + assertEquals(List.of("done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + } + + @Test + public void testHandleTask_generateFailedWithoutTerminal() { + // mock 方法：生成失败已回填，error 尚未写入 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO() + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent("\n\n> 本条回复生成失败")); + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of()); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true); + // 断言：按已回填的固定提示补写，不改为中断 + assertEquals(List.of("error", "done"), captureResultTypes()); + assertEquals(List.of("本条回复生成失败", ""), captureResultData()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + } + + @Test + public void testHandleTask_failedBeforeStart() { + // mock 方法：投递超时后补偿已标记失败，worker 才开始处理 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO() + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent("")); + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of()); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：不生成、不改写消息，补齐终态后确认 + verifyNoInteractions(agentService, llmChatTool, mcpToolFactory, skillToolFactory); + verify(sessionMessageService, never()).updateSessionMessage(any()); + assertEquals(List.of("error", "done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + } + + @Test + public void testHandleTask_injectMemory() { + // mock 方法 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L).setMemoryEnabled(true)); + when(sessionService.getSession(SESSION_ID)).thenReturn(new Ai1SessionDO().setId(SESSION_ID).setUserId(100L)); + when(memoryTool.loadMemory(AGENT_ID, 100L)).thenReturn("## 用户画像\n- 程序员"); + when(memoryTool.mergePrompt("你是 客服 的智能助手。", "## 用户画像\n- 程序员")).thenReturn("合并后的系统指令"); + when(llmChatTool.chat(any(), eq("合并后的系统指令"), anyList(), eq("你好"), anyList(), anyList(), eq(SESSION_ID), any(), any())) + .thenReturn(new Ai1LlmChatTool.ChatText("您好", "")); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言 + verify(sessionMessageService).updateSessionMessage(argThat(message -> + Ai1SessionMessageStatusEnum.SUCCESS.getStatus().equals(message.getStatus()))); + verify(skillToolFactory).buildTools(any(), same(shellTracker)); + } + + @Test + public void testHandleTask_injectMemoryFailed() { + // mock 方法：记忆查询失败时降级为原系统指令 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L).setMemoryEnabled(true)); + when(sessionService.getSession(SESSION_ID)).thenThrow(new IllegalStateException("数据库不可用")); + when(llmChatTool.chat(any(), eq("你是 客服 的智能助手。"), anyList(), eq("你好"), anyList(), anyList(), eq(SESSION_ID), any(), any())) + .thenReturn(new Ai1LlmChatTool.ChatText("您好", "")); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言 + verifyNoInteractions(memoryTool); + verify(sessionMessageService).updateSessionMessage(argThat(message -> + Ai1SessionMessageStatusEnum.SUCCESS.getStatus().equals(message.getStatus()))); + } + + @Test + public void testHandleTask_stoppedWhileQueued() { + // mock 排队期间已写入停止意图 + when(stopBucket.isExists()).thenReturn(true); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：不调用模型，人工终止只写 done，并清除中断标志 + verifyNoInteractions(agentService, llmChatTool); + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 本条回复已被人工终止".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + assertEquals(List.of("done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + assertFalse(Thread.currentThread().isInterrupted()); + } + + @Test + public void testHandleTask_stoppedDuringPreparation() { + // mock 方法：装配工具期间收到停止通知 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L)); + when(mcpToolFactory.buildTools(any())).thenAnswer(invocation -> { + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID); + return List.of(); + }); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：不再调用模型 + verifyNoInteractions(llmChatTool); + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 本条回复已被人工终止".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + assertEquals(List.of("ping", "done"), captureResultTypes()); + verify(shellTracker, atLeastOnce()).cancel(); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + assertFalse(Thread.currentThread().isInterrupted()); + } + + @Test + public void testHandleTask_stoppedDuringGenerate() { + // mock 方法：输出第一段后停止，之后的增量不再写入 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L)); + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any())) + .thenAnswer(invocation -> { + Consumer onContent = invocation.getArgument(8); + onContent.accept("第一段"); + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID); + onContent.accept("第二段"); + return new Ai1LlmChatTool.ChatText("第一段第二段", ""); + }); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：保留已生成内容，只提交一个终态 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段\n\n> 本条回复已被人工终止".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + assertEquals(List.of("ping", "message", "done"), captureResultTypes()); + assertFalse(Thread.currentThread().isInterrupted()); + } + + @Test + public void testStop_interruptRunning() { + // mock 方法：生成期间节点停机 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L)); + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any())) + .thenAnswer(invocation -> { + sessionStreamTool.stop(); + return new Ai1LlmChatTool.ChatText("部分", ""); + }); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：保留内容并标记中断，另写 error，确认后不会被重新生成 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "部分\n\n> 生成任务已中断，请重新发送".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + assertEquals(List.of("ping", "error", "done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + } + + @Test + public void testHandleTask_stoppedBeforeRegister() { + // mock 方法：查询消息后、登记上下文前节点停机 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenAnswer(invocation -> { + sessionStreamTool.stop(); + return new Ai1SessionMessageDO().setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + }); + + // 调用 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + // 断言：不调用模型、不装配工具，按中断收尾并确认 + verifyNoInteractions(agentService, llmChatTool, mcpToolFactory); + verify(stopBucket, never()).isExists(); + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 生成任务已中断，请重新发送".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + assertEquals(List.of("error", "done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + assertTrue(getRunContexts().isEmpty()); + assertFalse(Thread.currentThread().isInterrupted()); + } + + @Test + public void testConsumeLoop_stoppedAfterRead() { + // mock 方法：阻塞读取返回任务时，节点已停机 + when(taskStream.readGroup(anyString(), anyString(), any(StreamReadGroupArgs.class))).thenAnswer(invocation -> { + sessionStreamTool.stop(); + return Map.of(RECORD_ID, buildTaskFields()); + }); + + // 调用 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "consumeLoop", "worker-test"); + // 断言：不调用模型，按中断收尾后退出循环 + verifyNoInteractions(agentService, llmChatTool); + assertEquals(List.of("error", "done"), captureResultTypes()); + verify(taskStream).ack("ai1-session-workers", RECORD_ID); + assertFalse(Thread.currentThread().isInterrupted()); + } + + @Test + public void testCancel_queued() { + // mock 方法：会话中有一条生成中的助手消息 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of(MESSAGE_ID)); + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO().setId(MESSAGE_ID) + .setSessionId(SESSION_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus())); + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenReturn( + Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成"))); + when(taskLock.isHeldByCurrentThread()).thenReturn(true); + + // 调用 + sessionStreamTool.cancel(SESSION_ID); + // 断言：写入停止意图并广播 + verify(stopBucket).set("1", Duration.ofSeconds(3600)); + verify(stopTopic).publish(String.valueOf(MESSAGE_ID)); + verify(redissonClient, times(1)).getBucket(anyString(), any(Codec.class)); + // 断言：没有 worker 执行，直接收尾 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "已生成\n\n> 本条回复已被人工终止".equals(message.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus()))); + verify(sessionService).touchSession(SESSION_ID); + assertEquals(List.of("done"), captureResultTypes()); + verify(taskLock).unlock(); + verify(sessionMessageService, never()).getSessionMessageListBySessionId(any()); + } + + @Test + public void testCancel_running() { + // mock 方法：任务正在其他 worker 执行 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of(MESSAGE_ID)); + when(taskLock.tryLock()).thenReturn(false); + + // 调用 + sessionStreamTool.cancel(SESSION_ID); + // 断言：只写意图与通知，由执行节点收尾 + verify(stopBucket).set("1", Duration.ofSeconds(3600)); + verify(stopTopic).publish(String.valueOf(MESSAGE_ID)); + verify(sessionMessageService, never()).updateSessionMessage(any()); + verify(resultStream, never()).add(any()); + } + + @Test + public void testCancel_noGenerating() { + // mock 方法 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of()); + + // 调用 + sessionStreamTool.cancel(SESSION_ID); + // 断言：没有生成中的消息时无副作用 + verifyNoInteractions(stopBucket, stopTopic, taskLock); + } + + @Test + @Timeout(30) + public void testCheckStopRequests() throws Exception { + // 准备参数 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(new Thread(() -> {}), shellTracker, getReapers()); + Map runContexts = getRunContexts(); + runContexts.put(MESSAGE_ID, context); + // mock 订阅断开期间写入的停止意图；进程回收一直阻塞 + when(stopBucket.isExists()).thenReturn(true); + CountDownLatch killing = new CountDownLatch(1); + CountDownLatch releaseKill = new CountDownLatch(1); + doAnswer(invocation -> { + killing.countDown(); + releaseKill.await(10, TimeUnit.SECONDS); + return null; + }).when(shellTracker).cancel(); + + try { + // 调用：回收阻塞时检查线程仍立即返回 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "checkStopRequests"); + // 断言：回收在独立线程执行 + assertTrue(killing.await(5, TimeUnit.SECONDS)); + assertEquals(Ai1SessionStreamTool.TERMINATED_NOTICE, context.finish()); + } finally { + releaseKill.countDown(); + awaitReapers(); + } + } + + @Test + @Timeout(30) + public void testStop_slowReapNotBlockOthers() throws Exception { + // 准备参数：两个执行中的生成，第一个的进程回收一直阻塞 + Ai1SkillShellTracker slowTracker = mock(Ai1SkillShellTracker.class); + Ai1SkillShellTracker fastTracker = mock(Ai1SkillShellTracker.class); + CountDownLatch killing = new CountDownLatch(1); + CountDownLatch releaseKill = new CountDownLatch(1); + CountDownLatch fastReaped = new CountDownLatch(1); + doAnswer(invocation -> { + killing.countDown(); + releaseKill.await(10, TimeUnit.SECONDS); + return null; + }).when(slowTracker).cancel(); + doAnswer(invocation -> { + fastReaped.countDown(); + return null; + }).when(fastTracker).cancel(); + CountDownLatch releaseWorker = new CountDownLatch(1); + CountDownLatch interrupted = new CountDownLatch(2); + Thread worker1 = startWorker(releaseWorker, interrupted); + Thread worker2 = startWorker(releaseWorker, interrupted); + getRunContexts().put(1L, new Ai1SessionStreamTool.RunContext(worker1, slowTracker, getReapers())); + getRunContexts().put(2L, new Ai1SessionStreamTool.RunContext(worker2, fastTracker, getReapers())); + Thread stopping = new Thread(sessionStreamTool::stop); + + try { + // 调用 + stopping.start(); + // 断言：慢回收进行中，两个生成都被中断，另一个回收不受影响；stop 等待回收结束 + assertTrue(killing.await(5, TimeUnit.SECONDS)); + assertTrue(interrupted.await(5, TimeUnit.SECONDS)); + assertTrue(fastReaped.await(5, TimeUnit.SECONDS)); + assertTrue(stopping.isAlive()); + } finally { + releaseKill.countDown(); + releaseWorker.countDown(); + stopping.join(10_000); + worker1.join(5000); + worker2.join(5000); + } + assertFalse(stopping.isAlive()); + } + + @Test + @Timeout(30) + @EnabledOnOs({OS.MAC, OS.LINUX}) + public void testRunContext_cancelReapsWhileWorkerIgnoresInterrupt(@TempDir Path workDir) throws Exception { + // 准备参数：worker 吞掉中断，模拟模型或工具库不响应中断；Shell 写入自身 PID 后 exec 为 sleep + CountDownLatch releaseWorker = new CountDownLatch(1); + Thread worker = new Thread(() -> { + while (releaseWorker.getCount() > 0) { + try { + releaseWorker.await(); + } catch (InterruptedException ignored) { + // 忽略中断，继续等待 + } + } + }); + worker.start(); + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(workDir).build()); + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(worker, tracker, getReapers()); + ExecHandle handle = tracker.start(ExecSpec.of("echo $$ > pid.tmp && mv pid.tmp pid && exec sleep 60")); + ProcessHandle process = null; + try { + process = awaitDescendant(workDir.resolve("pid")); + assertTrue(process.isAlive()); + + // 调用 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE); + // 断言：worker 仍未返回，在册的后台进程已被回收 + process.onExit().get(10, TimeUnit.SECONDS); + assertFalse(process.isAlive()); + assertFalse(handle.isAlive()); + assertTrue(worker.isAlive()); + } finally { + // 只清理本用例启动的线程、进程 + releaseWorker.countDown(); + worker.join(5000); + awaitReapers(); + if (process != null && process.isAlive()) { + process.destroyForcibly(); + } + if (handle.isAlive()) { + handle.kill(); + } + } + } + + @Test + public void testRunContext_cancelThenFinish() throws Exception { + // 准备参数 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(Thread.currentThread(), shellTracker, getReapers()); + + // 调用：重复停止只记录第一次 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE); + context.cancel(Ai1SessionStreamTool.INTERRUPTED_MESSAGE); + // 断言 + assertTrue(Thread.currentThread().isInterrupted()); + assertThrows(CancellationException.class, context::checkActive); + assertEquals(Ai1SessionStreamTool.TERMINATED_NOTICE, context.finish()); + assertFalse(Thread.currentThread().isInterrupted()); + // 断言：只启动一次回收 + awaitReapers(); + verify(shellTracker, times(1)).cancel(); + } + + @Test + public void testRunContext_lateCancel() { + // 准备参数 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(Thread.currentThread(), shellTracker, getReapers()); + assertNull(context.finish()); + + // 调用：结束后迟到的停止 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE); + // 断言：不中断 worker 处理下一个任务，也不启动回收 + assertFalse(Thread.currentThread().isInterrupted()); + assertTrue(getReapers().isEmpty()); + verifyNoInteractions(shellTracker); + } + + @Test + @Timeout(30) + public void testStop_waitReaperAfterWorkerExit() throws Exception { + // mock 方法：生成中启动后台进程，人工停止后进程回收一直阻塞 + ExecBackend backend = mock(ExecBackend.class); + ExecHandle handle = mock(ExecHandle.class); + ExecSpec spec = ExecSpec.of("sleep 60"); + when(backend.start(spec)).thenReturn(handle); + CountDownLatch killing = new CountDownLatch(1); + CountDownLatch releaseKill = new CountDownLatch(1); + CountDownLatch killed = new CountDownLatch(1); + doAnswer(invocation -> { + killing.countDown(); + releaseKill.await(10, TimeUnit.SECONDS); + killed.countDown(); + return null; + }).when(handle).kill(); + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(backend); + when(skillToolFactory.newShellTracker(AGENT_ID)).thenReturn(tracker); + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L)); + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any())) + .thenAnswer(invocation -> { + tracker.start(spec); + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID); + // 回收线程认领进程后再返回；worker 线程已被中断，用自旋等待 + long deadline = System.currentTimeMillis() + 5000; + while (killing.getCount() > 0 && System.currentTimeMillis() < deadline) { + Thread.onSpinWait(); + } + return new Ai1LlmChatTool.ChatText("部分", ""); + }); + Thread stopping = new Thread(sessionStreamTool::stop); + + try { + // 调用：worker 结束并移除上下文，回收仍阻塞 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields()); + assertTrue(killing.await(5, TimeUnit.SECONDS)); + assertTrue(getRunContexts().isEmpty()); + assertEquals(1, killed.getCount()); + // 调用：节点停止 + stopping.start(); + // 断言：stop 仍等待已结束上下文的回收 + stopping.join(1000); + assertTrue(stopping.isAlive()); + } finally { + releaseKill.countDown(); + stopping.join(10_000); + awaitReapers(); + } + // 断言：回收完成后 stop 返回，worker 未重复回收 + assertFalse(stopping.isAlive()); + assertEquals(0, killed.getCount()); + verify(handle, times(1)).kill(); + assertTrue(getReapers().isEmpty()); + } + + @Test + public void testConsumeLoop_noGroup() { + // mock 方法：任务队列被删除后消费组丢失，重建后退出循环 + ReflectionTestUtils.setField(sessionStreamTool, "running", true); + when(taskStream.readGroup(anyString(), anyString(), any(StreamReadGroupArgs.class))) + .thenThrow(new RedisException("NOGROUP No such key 'ai1:session:tasks' or consumer group")); + doAnswer(invocation -> { + ReflectionTestUtils.setField(sessionStreamTool, "running", false); + return null; + }).when(taskStream).createGroup(any(StreamCreateGroupArgs.class)); + + // 调用 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "consumeLoop", "worker-test"); + // 断言 + verify(taskStream).createGroup(any(StreamCreateGroupArgs.class)); + } + + @Test + public void testFailStaleMessage_notStale() { + // 调用，并断言：未超时时不加锁，避免 worker 误判为在途 + assertFalse(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now()))); + verifyNoInteractions(redissonClient, sessionMessageService); + } + + @Test + public void testFailStaleMessage_running() { + // mock 方法：仍有 worker 持有执行锁 + when(taskLock.tryLock()).thenReturn(false); + + // 调用，并断言 + assertFalse(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10)))); + verify(sessionMessageService, never()).getSessionMessage(any()); + verify(sessionMessageService, never()).updateSessionMessage(any()); + verify(resultStream, never()).add(any()); + } + + @Test + public void testFailStaleMessage_success() { + // mock 方法：任务记录已裁剪，没有 worker 执行 + Ai1SessionMessageDO message = buildGeneratingMessage(LocalDateTime.now().minusMinutes(10)); + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(message); + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenReturn( + Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成"))); + + // 调用，并断言：与接管任务一致，保留部分结果并标记中断 + assertTrue(sessionStreamTool.failStaleMessage(message)); + verify(sessionMessageService).updateSessionMessage(argThat(updateObj -> "已生成\n\n> 生成任务已中断，请重新发送".equals(updateObj.getContent()) + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(updateObj.getStatus()))); + verify(sessionService).touchSession(SESSION_ID); + assertEquals(List.of("error", "done"), captureResultTypes()); + verify(taskLock).unlock(); + verifyNoInteractions(llmChatTool); + } + + @Test + public void testFailStaleMessage_completedConcurrently() { + // mock 方法：加锁前 worker 已完成 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10)) + .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus())); + + // 调用，并断言：不再改写，归档可继续 + assertTrue(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10)))); + verify(sessionMessageService, never()).updateSessionMessage(any()); + verify(resultStream, never()).add(any()); + verify(taskLock).unlock(); + } + + @Test + public void testFailStaleMessage_lockLostDuringPartialRead() { + // mock 方法：读取部分结果期间执行锁失效，已被新持有者接管 + Ai1SessionMessageDO message = buildGeneratingMessage(LocalDateTime.now().minusMinutes(10)); + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(message); + AtomicBoolean owned = new AtomicBoolean(true); + when(taskLock.isHeldByThread(anyLong())).thenAnswer(invocation -> owned.get()); + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenAnswer(invocation -> { + owned.set(false); + return Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成")); + }); + + // 调用，并断言：不改动消息、结果流，也不释放新持有者的锁 + assertThrows(IllegalStateException.class, () -> sessionStreamTool.failStaleMessage(message)); + verify(sessionMessageService, never()).updateSessionMessage(any()); + verify(sessionService, never()).touchSession(any()); + verify(resultStream, never()).add(any()); + verify(taskLock, never()).unlock(); + } + + // ========== mock 方法 ========== + + private void mockGenerate(Ai1AgentDO agent) { + when(agentService.validateAgentExists(AGENT_ID)).thenReturn(agent); + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType())); + when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt())) + .thenReturn(new ArrayList<>()); + } + + private static Thread startWorker(CountDownLatch release, CountDownLatch interrupted) { + Thread worker = new Thread(() -> { + try { + release.await(); + } catch (InterruptedException e) { + interrupted.countDown(); + } + }); + worker.start(); + return worker; + } + // ========== 测试数据 ========== private static Map buildTaskFields() { @@ -449,6 +1058,33 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { return fields; } + private static Ai1SessionMessageDO buildGeneratingMessage(LocalDateTime createTime) { + Ai1SessionMessageDO message = new Ai1SessionMessageDO().setId(MESSAGE_ID).setSessionId(SESSION_ID) + .setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + message.setCreateTime(createTime); + return message; + } + + /** + * 等待 PID 文件写入，且该进程属于当前 JVM 的子孙进程 + */ + private static ProcessHandle awaitDescendant(Path pidFile) throws Exception { + long deadline = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadline) { + if (Files.exists(pidFile)) { + long pid = Long.parseLong(Files.readString(pidFile).trim()); + Optional process = ProcessHandle.current().descendants() + .filter(descendant -> descendant.pid() == pid).findFirst(); + if (process.isPresent()) { + return process.get(); + } + } + Thread.sleep(50); + } + return fail("未找到本用例启动的 Shell 进程"); + } + // ========== 断言工具 ========== @SuppressWarnings({"unchecked", "rawtypes"}) @@ -463,6 +1099,36 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest { return types; } + @SuppressWarnings({"unchecked", "rawtypes"}) + private List captureResultData() { + ArgumentCaptor captor = ArgumentCaptor.forClass(StreamAddArgs.class); + verify(resultStream, atLeastOnce()).add(captor.capture()); + List data = new ArrayList<>(); + for (StreamAddArgs args : captor.getAllValues()) { + data.add(String.valueOf(readEntries(args).get("data"))); + } + return data; + } + + @SuppressWarnings("unchecked") + private Set getReapers() { + return (Set) ReflectionTestUtils.getField(sessionStreamTool, "reapers"); + } + + /** + * 等待本用例启动的回收线程结束 + */ + private void awaitReapers() throws InterruptedException { + for (Thread reaper : getReapers()) { + reaper.join(10_000); + } + } + + @SuppressWarnings("unchecked") + private Map getRunContexts() { + return (Map) ReflectionTestUtils.getField(sessionStreamTool, "runContexts"); + } + @SuppressWarnings("unchecked") private static Map readEntries(StreamAddArgs args) { try { diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionTaskRecoveryIntegrationTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionTaskRecoveryIntegrationTest.java index 136e5aa..e0bbfff 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionTaskRecoveryIntegrationTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/session/Ai1SessionTaskRecoveryIntegrationTest.java @@ -9,7 +9,9 @@ import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties; import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool; import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory; +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool; import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool; +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker; import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory; import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; import cn.iocoder.yudao.module.ai1.service.knowledge.Ai1KnowledgeBaseService; @@ -26,6 +28,7 @@ import org.mockito.Mock; import org.redisson.Redisson; import org.redisson.api.RScript; import org.redisson.api.RStream; +import org.redisson.api.RTopic; import org.redisson.api.RedissonClient; import org.redisson.api.stream.StreamAddArgs; import org.redisson.api.stream.StreamCreateGroupArgs; @@ -87,6 +90,10 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { private Ai1SkillToolFactory skillToolFactory; @Mock private Ai1RagTool ragTool; + @Mock + private Ai1MemoryTool memoryTool; + @Mock + private Ai1SkillShellTracker shellTracker; private RedissonClient firstClient; private RedissonClient secondClient; @@ -114,18 +121,21 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { message.set(new Ai1SessionMessageDO().setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()) .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus())); - // mock 业务方法 + // mock 业务方法（排队中停止的用例不调用模型） when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenAnswer(invocation -> message.get()); + lenient().when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenAnswer(invocation -> + Ai1SessionMessageStatusEnum.isGenerating(message.get().getStatus()) ? List.of(MESSAGE_ID) : List.of()); doAnswer(invocation -> { message.set(invocation.getArgument(0)); return null; }).when(sessionMessageService).updateSessionMessage(any()); - when(agentService.validateAgentExists(1L)).thenReturn(new Ai1AgentDO().setId(1L).setName("客服") + lenient().when(skillToolFactory.newShellTracker(1L)).thenReturn(shellTracker); + lenient().when(agentService.validateAgentExists(1L)).thenReturn(new Ai1AgentDO().setId(1L).setName("客服") .setProviderId(1L).setModelId(2L)); - when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType())); - when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt())) + lenient().when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType())); + lenient().when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt())) .thenReturn(new ArrayList<>()); - when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any())) + lenient().when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any())) .thenAnswer(invocation -> { Consumer onContent = invocation.getArgument(8); onContent.accept("第一段"); @@ -143,6 +153,7 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); taskStream.delete(); resultStream.delete(); + secondClient.getBucket(String.format(SESSION_STOP, MESSAGE_ID)).delete(); secondClient.getLock(String.format(SESSION_TASK_LOCK, MESSAGE_ID)).forceUnlock(); firstClient.shutdown(); secondClient.shutdown(); @@ -195,7 +206,7 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { secondWorker.handleTask(recordId, reclaimed.get(recordId), true); // 断言接管只保留部分内容并标记失败，不再次调用模型 - assertEquals("第一段", message.get().getContent()); + assertEquals("第一段\n\n> 生成任务已中断，请重新发送", message.get().getContent()); assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus()); assertEquals(List.of("ping", "message", "error", "done"), resultTypes()); assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal()); @@ -204,10 +215,108 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { // 调用旧节点恢复，断言不追加旧增量或覆盖回填结果 release.countDown(); generation.get(5, TimeUnit.SECONDS); - assertEquals("第一段", message.get().getContent()); + assertEquals("第一段\n\n> 生成任务已中断，请重新发送", message.get().getContent()); assertEquals(List.of("ping", "message", "error", "done"), resultTypes()); } + @Test + public void testStop_crossNode() throws Exception { + // 准备任务，交给第一个节点生成；第一个节点订阅停止通知 + startWithoutWorker(firstWorker); + StreamMessageId recordId = addTask(); + Map fields = taskStream.readGroup(SESSION_TASK_GROUP, "first", + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId); + Future generation = executor.submit(() -> firstWorker.handleTask(recordId, fields)); + assertTrue(generated.await(5, TimeUnit.SECONDS)); + + try { + // 调用：请求落在第二个节点 + secondWorker.cancel(SESSION_ID); + generation.get(5, TimeUnit.SECONDS); + + // 断言：保留已生成内容，人工终止只写 done，并确认任务 + assertEquals("第一段\n\n> 本条回复已被人工终止", message.get().getContent()); + assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus()); + assertEquals(List.of("ping", "message", "done"), resultTypes()); + assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal()); + verify(shellTracker, atLeastOnce()).cancel(); + // 断言：worker 线程的中断标志已清除，不影响下一轮 + assertFalse(executor.submit(() -> Thread.currentThread().isInterrupted()).get(5, TimeUnit.SECONDS)); + } finally { + firstWorker.stop(); + } + } + + @Test + public void testStop_intentWithoutSubscription() throws Exception { + // 准备任务，交给第一个节点生成；模拟停止通知的订阅已断开 + startWithoutWorker(firstWorker); + RTopic stopTopic = firstClient.getTopic(SESSION_STOP_TOPIC, StringCodec.INSTANCE); + stopTopic.removeListener((Integer) ReflectionTestUtils.getField(firstWorker, "stopListenerId")); + StreamMessageId recordId = addTask(); + Map fields = taskStream.readGroup(SESSION_TASK_GROUP, "first", + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId); + Future generation = executor.submit(() -> firstWorker.handleTask(recordId, fields)); + assertTrue(generated.await(5, TimeUnit.SECONDS)); + + try { + // 调用：停止通知丢失，由定时检查停止意图兜底 + secondWorker.cancel(SESSION_ID); + generation.get(5, TimeUnit.SECONDS); + + // 断言 + assertEquals("第一段\n\n> 本条回复已被人工终止", message.get().getContent()); + assertEquals(List.of("ping", "message", "done"), resultTypes()); + } finally { + firstWorker.stop(); + } + } + + @Test + public void testStop_queued() { + // 准备任务，尚未被任何 worker 读取 + StreamMessageId recordId = addTask(); + + // 调用：无人执行时由请求节点直接收尾 + secondWorker.cancel(SESSION_ID); + // 断言 + assertEquals("\n\n> 本条回复已被人工终止", message.get().getContent()); + assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus()); + assertEquals(List.of("done"), resultTypes()); + + // 调用：worker 之后读到该任务 + Map fields = taskStream.readGroup(SESSION_TASK_GROUP, "first", + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId); + firstWorker.handleTask(recordId, fields); + // 断言：不调用模型，补齐终态后确认 + verifyNoInteractions(llmChatTool); + assertEquals(List.of("done"), resultTypes()); + assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal()); + } + + @Test + public void testConsume_missingGroupRecovered() throws Exception { + // 准备：启动一个 worker 后删除任务队列，消费组随之丢失 + YudaoAi1Properties properties = (YudaoAi1Properties) ReflectionTestUtils.getField(secondWorker, "ai1Properties"); + properties.getSession().getWorker().setCount(1); + secondWorker.start(); + try { + taskStream.delete(); + addTask(); + release.countDown(); + + // 断言：worker 重建消费组后继续消费 + long deadline = System.currentTimeMillis() + 20_000; + while (Ai1SessionMessageStatusEnum.isGenerating(message.get().getStatus()) && System.currentTimeMillis() < deadline) { + Thread.sleep(200); + } + assertEquals(Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message.get().getStatus()); + assertEquals("第一段第二段", message.get().getContent()); + } finally { + secondWorker.stop(); + } + } + // ========== 测试数据 ========== private RedissonClient createClient() { @@ -241,9 +350,21 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest { ReflectionTestUtils.setField(worker, "mcpToolFactory", mcpToolFactory); ReflectionTestUtils.setField(worker, "skillToolFactory", skillToolFactory); ReflectionTestUtils.setField(worker, "ragTool", ragTool); + ReflectionTestUtils.setField(worker, "memoryTool", memoryTool); + // 直接调用 handleTask 的用例按运行中处理，与 start 之后一致 + ReflectionTestUtils.setField(worker, "running", true); return worker; } + /** + * 启动节点但不拉起消费线程，只订阅停止通知、检查停止意图 + */ + private static void startWithoutWorker(Ai1SessionStreamTool worker) { + YudaoAi1Properties properties = (YudaoAi1Properties) ReflectionTestUtils.getField(worker, "ai1Properties"); + properties.getSession().getWorker().setCount(0); + worker.start(); + } + private StreamMessageId addTask() { return taskStream.add(StreamAddArgs.entries(Map.of("messageId", String.valueOf(MESSAGE_ID), "agentId", "1", "sessionId", String.valueOf(SESSION_ID), "content", "你好"))); diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java new file mode 100644 index 0000000..9fce0d0 --- /dev/null +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java @@ -0,0 +1,175 @@ +package cn.iocoder.yudao.module.ai1.harness.skill; + +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; +import org.junit.jupiter.api.condition.EnabledOnOs; +import org.junit.jupiter.api.condition.OS; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.Mock; +import org.mockito.stubbing.Answer; +import org.springaicommunity.agent.common.exec.ExecBackend; +import org.springaicommunity.agent.common.exec.ExecHandle; +import org.springaicommunity.agent.common.exec.ExecResult; +import org.springaicommunity.agent.common.exec.ExecSpec; +import org.springaicommunity.agent.exec.LocalExecBackend; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.*; + +/** + * {@link Ai1SkillShellTracker} 的单元测试 + * + * @author 芋道源码 + */ +public class Ai1SkillShellTrackerTest extends BaseMockitoUnitTest { + + private static final ExecSpec SPEC = ExecSpec.of("sleep 60"); + + @Mock + private ExecBackend delegate; + @Mock + private ExecHandle handle1; + @Mock + private ExecHandle handle2; + + @Test + public void testRun_delegate() { + // mock 方法 + ExecResult result = ExecResult.completed(0, "ok", ""); + when(delegate.run(SPEC)).thenReturn(result); + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate); + + // 调用，并断言：同步命令自带超时，不登记 + assertSame(result, tracker.run(SPEC)); + tracker.cancel(); + verifyNoInteractions(handle1); + } + + @Test + public void testCancel_killAll() { + // mock 方法：第一个进程回收失败，不影响其他进程 + when(delegate.start(SPEC)).thenReturn(handle1, handle2); + doThrow(new IllegalStateException("进程已退出")).when(handle1).kill(); + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate); + tracker.start(SPEC); + tracker.start(SPEC); + + // 调用：重复终止幂等 + tracker.cancel(); + tracker.cancel(); + // 断言 + verify(handle1).kill(); + verify(handle2).kill(); + } + + @Test + @Timeout(30) + public void testCancel_concurrentKillOnce() throws Exception { + // mock 方法：第一个进程回收期间，另一个线程（如 worker 结束）并发回收 + when(delegate.start(SPEC)).thenReturn(handle1, handle2); + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate); + tracker.start(SPEC); + tracker.start(SPEC); + AtomicBoolean first = new AtomicBoolean(true); + Answer concurrentCancel = invocation -> { + if (first.getAndSet(false)) { + Thread other = new Thread(tracker::cancel); + other.start(); + other.join(5000); + } + return null; + }; + doAnswer(concurrentCancel).when(handle1).kill(); + doAnswer(concurrentCancel).when(handle2).kill(); + + // 调用 + tracker.cancel(); + // 断言：每个进程只回收一次 + verify(handle1, times(1)).kill(); + verify(handle2, times(1)).kill(); + } + + @Test + public void testStart_afterCancel() { + // 准备参数 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate); + tracker.cancel(); + + // 调用，并断言异常：终止后拒绝再启动 + assertThrows(IllegalStateException.class, () -> tracker.start(SPEC)); + verify(delegate, never()).start(any()); + } + + @Test + public void testStart_cancelledDuringStart() { + // mock 方法：进程启动期间被终止 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate); + when(delegate.start(SPEC)).thenAnswer(invocation -> { + tracker.cancel(); + return handle1; + }); + + // 调用 + assertSame(handle1, tracker.start(SPEC)); + // 断言：登记后发现已终止，立即回收 + verify(handle1).kill(); + } + + @Test + @Timeout(30) + @EnabledOnOs({OS.MAC, OS.LINUX}) + public void testCancel_realProcess(@TempDir Path workDir) throws Exception { + // 准备参数：Shell 写入自身 PID 后 exec 为 sleep，PID 不变 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(workDir).build()); + ExecHandle handle = tracker.start(ExecSpec.of("echo $$ > pid.tmp && mv pid.tmp pid && exec sleep 60")); + ProcessHandle process = null; + try { + process = awaitDescendant(workDir.resolve("pid")); + assertTrue(process.isAlive()); + + // 调用 + tracker.cancel(); + // 断言 + process.onExit().get(10, TimeUnit.SECONDS); + assertFalse(process.isAlive()); + assertFalse(handle.isAlive()); + } finally { + // 只清理本用例启动的进程 + if (process != null && process.isAlive()) { + process.destroyForcibly(); + } + if (handle.isAlive()) { + handle.kill(); + } + } + } + + // ========== 测试数据 ========== + + /** + * 等待 PID 文件写入，且该进程属于当前 JVM 的子孙进程 + */ + private static ProcessHandle awaitDescendant(Path pidFile) throws Exception { + long deadline = System.currentTimeMillis() + 5_000; + while (System.currentTimeMillis() < deadline) { + if (Files.exists(pidFile)) { + long pid = Long.parseLong(Files.readString(pidFile).trim()); + Optional process = ProcessHandle.current().descendants() + .filter(descendant -> descendant.pid() == pid).findFirst(); + if (process.isPresent()) { + return process.get(); + } + } + Thread.sleep(50); + } + return fail("未找到本用例启动的 Shell 进程"); + } + +} diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactoryTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactoryTest.java index c7fcc5f..8725f29 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactoryTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactoryTest.java @@ -16,6 +16,8 @@ import org.junit.jupiter.api.io.TempDir; import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.Spy; +import org.springaicommunity.agent.common.exec.ExecResult; +import org.springaicommunity.agent.common.exec.ExecSpec; import org.springframework.ai.tool.ToolCallback; import java.nio.file.Path; @@ -64,7 +66,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { mockSkillFiles("echo hi"); // 调用 - List tools = skillToolFactory.buildTools(buildAgent()); + List tools = buildTools(buildAgent()); // 断言 assertEquals(6, tools.size()); assertInstanceOf(ToolCallback.class, tools.get(0)); @@ -82,8 +84,8 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { Ai1AgentDO agent = buildAgent(); // 调用 - skillToolFactory.buildTools(agent); - skillToolFactory.buildTools(agent); + buildTools(agent); + buildTools(agent); // 断言 verify(skillFileService, times(1)).getSkillFileListBySkillId(SKILL_ID); @@ -94,7 +96,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { mockSkillFiles("echo v2"); // 调用 - skillToolFactory.buildTools(agent); + buildTools(agent); // 断言 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID); assertEquals("echo v2", FileUtil.readUtf8String(agentRoot().resolve("pdf-reader/scripts/run.sh").toFile())); @@ -107,13 +109,13 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0)); mockSkillFiles("echo hi"); Ai1AgentDO agent = buildAgent(); - skillToolFactory.buildTools(agent); + buildTools(agent); // mock 物化目录被删除 FileUtil.del(agentRoot()); // 调用 - skillToolFactory.buildTools(agent); + buildTools(agent); // 断言 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID); assertTrue(FileUtil.exist(agentRoot().resolve("pdf-reader/SKILL.md").toFile())); @@ -125,7 +127,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { mockSkill(CommonStatusEnum.DISABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0)); // 调用 - List tools = skillToolFactory.buildTools(buildAgent()); + List tools = buildTools(buildAgent()); // 断言只保留执行工具，不物化 SKILL 文件 assertEquals(5, tools.size()); assertTrue(tools.stream().noneMatch(ToolCallback.class::isInstance)); @@ -136,7 +138,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { @Test public void testBuildTools_noSkill() { // 调用，并断言 - assertTrue(skillToolFactory.buildTools(new Ai1AgentDO().setId(AGENT_ID).setSkillIds(List.of())).isEmpty()); + assertTrue(buildTools(new Ai1AgentDO().setId(AGENT_ID).setSkillIds(List.of())).isEmpty()); verify(skillService, never()).getSkillList(anyCollection()); } @@ -146,7 +148,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0)); mockSkillFiles("echo hi"); Ai1AgentDO agent = buildAgent(); - skillToolFactory.buildTools(agent); + buildTools(agent); assertTrue(FileUtil.exist(agentRoot().toFile())); // 调用 @@ -155,7 +157,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { assertFalse(FileUtil.exist(agentRoot().toFile())); // 调用 - skillToolFactory.buildTools(agent); + buildTools(agent); // 断言重新物化 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID); } @@ -176,7 +178,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { } // 调用 - ToolCallback tool = (ToolCallback) skillToolFactory.buildTools(agent).get(0); + ToolCallback tool = (ToolCallback) buildTools(agent).get(0); String result = tool.call("{\"command\":\"same\"}"); // 断言 assertTrue(result.contains("first-content")); @@ -185,14 +187,32 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest { // 调用 skillToolFactory.evict(AGENT_ID); - ToolCallback rebuiltTool = (ToolCallback) skillToolFactory.buildTools(agent).get(0); + ToolCallback rebuiltTool = (ToolCallback) buildTools(agent).get(0); String rebuiltResult = rebuiltTool.call("{\"command\":\"same\"}"); // 断言 assertEquals(result, rebuiltResult); } + @Test + public void testNewShellTracker_workingDirectory() throws Exception { + // mock skillService 和 skillFileService 的方法 + mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0)); + mockSkillFiles("echo hi"); + buildTools(buildAgent()); + + // 调用 + ExecResult result = skillToolFactory.newShellTracker(AGENT_ID).run(ExecSpec.of("pwd")); + // 断言：Shell 命令限定在 Agent 物化根目录执行 + assertEquals(0, result.exitCode()); + assertEquals(agentRoot().toRealPath(), Path.of(result.stdout().trim()).toRealPath()); + } + // ========== 测试数据 ========== + private List buildTools(Ai1AgentDO agent) { + return skillToolFactory.buildTools(agent, skillToolFactory.newShellTracker(agent.getId())); + } + private Path agentRoot() { return tempDir.resolve("agent_" + AGENT_ID); } diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImplTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImplTest.java index 537a009..35253f8 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImplTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentServiceImplTest.java @@ -7,6 +7,7 @@ import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; import cn.iocoder.yudao.module.ai1.dal.mysql.agent.Ai1AgentMapper; import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum; import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory; +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService; import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService; import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService; import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO; @@ -45,6 +46,8 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest { @MockitoBean private Ai1SessionService sessionService; @MockitoBean + private Ai1MemoryService memoryService; + @MockitoBean private Ai1SkillToolFactory skillToolFactory; @Test @@ -123,6 +126,23 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest { } @Test + public void testGetAgentListByStatus_all() { + // mock 数据：开启、关闭、已删除的 Agent 各一个 + mockChatModel(); + Long enabledId = agentService.createAgent(buildAgentSaveReqVO()); + agentService.updateAgentStatus(enabledId, CommonStatusEnum.ENABLE.getStatus()); + Long disabledId = agentService.createAgent(buildAgentSaveReqVO()); + agentMapper.deleteById(agentService.createAgent(buildAgentSaveReqVO())); + + // 调用：状态为空时返回全部状态 + List list = agentService.getAgentListByStatus(null); + // 断言：按编号升序，不含已删除的 Agent + assertEquals(List.of(enabledId, disabledId), list.stream().map(Ai1AgentDO::getId).toList()); + assertEquals(CommonStatusEnum.ENABLE.getStatus(), list.get(0).getStatus()); + assertEquals(CommonStatusEnum.DISABLE.getStatus(), list.get(1).getStatus()); + } + + @Test public void testDeleteAgent_cascade() { // mock 数据 mockChatModel(); @@ -133,10 +153,29 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest { // 断言 assertNull(agentMapper.selectById(id)); verify(sessionService).deleteSessionListByAgentIds(Collections.singletonList(id)); + verify(memoryService).deleteMemoryListByAgentIds(Collections.singletonList(id)); verify(skillToolFactory).evict(id); } @Test + public void testCreateAgent_memoryEnabled() { + // mock 数据 + mockChatModel(); + + // 调用 + Long defaultId = agentService.createAgent(buildAgentSaveReqVO()); + Long enabledId = agentService.createAgent(buildAgentSaveReqVO().setMemoryEnabled(true)); + // 断言：默认关闭 + assertFalse(agentMapper.selectById(defaultId).getMemoryEnabled()); + assertTrue(agentMapper.selectById(enabledId).getMemoryEnabled()); + + // 调用：关闭 + agentService.updateAgent(buildAgentSaveReqVO().setId(enabledId).setMemoryEnabled(false)); + // 断言 + assertFalse(agentMapper.selectById(enabledId).getMemoryEnabled()); + } + + @Test public void testGetAgentCountByKnowledgeBaseIdsAndSkillIds() { // mock 数据 mockChatModel(); diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java new file mode 100644 index 0000000..a67274a --- /dev/null +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java @@ -0,0 +1,224 @@ +package cn.iocoder.yudao.module.ai1.service.memory; + +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum; +import cn.iocoder.yudao.framework.common.pojo.PageResult; +import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO; +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO; +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO; +import cn.iocoder.yudao.module.ai1.dal.mysql.memory.Ai1MemoryMapper; +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; +import jakarta.annotation.Resource; +import org.junit.jupiter.api.Test; +import org.springframework.context.annotation.Import; +import org.springframework.test.context.bean.override.mockito.MockitoBean; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException; +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MEMORY_NOT_EXISTS; +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.Mockito.when; + +/** + * {@link Ai1MemoryServiceImpl} 的单元测试 + * + * @author 芋道源码 + */ +@Import(Ai1MemoryServiceImpl.class) +public class Ai1MemoryServiceImplTest extends BaseDbUnitTest { + + @Resource + private Ai1MemoryServiceImpl memoryService; + + @Resource + private Ai1MemoryMapper memoryMapper; + + @MockitoBean + private Ai1AgentService agentService; + + @Test + public void testUpdateMemory_success() { + // mock 数据 + Ai1MemoryDO memory = insertMemory(1L, 100L); + + // 调用 + memoryService.updateMemory(new Ai1MemoryUpdateReqVO().setId(memory.getId()).setContent("## 用户画像\n- 新内容")); + // 断言：正文更新，水位线不变 + Ai1MemoryDO result = memoryMapper.selectById(memory.getId()); + assertEquals("## 用户画像\n- 新内容", result.getContent()); + assertEquals(10L, result.getLastMessageId()); + } + + @Test + public void testUpdateMemory_notExists() { + // 调用，并断言异常 + assertServiceException(() -> memoryService.updateMemory(new Ai1MemoryUpdateReqVO().setId(1L).setContent("内容")), + MEMORY_NOT_EXISTS); + } + + @Test + public void testDeleteMemoryListByIds_success() { + // mock 数据 + Ai1MemoryDO memory1 = insertMemory(1L, 100L); + Ai1MemoryDO memory2 = insertMemory(1L, 200L); + Ai1MemoryDO memory3 = insertMemory(2L, 100L); + + // 调用 + memoryService.deleteMemoryListByIds(Arrays.asList(memory1.getId(), memory2.getId())); + // 断言：逻辑删除后不再可见，同一用户可以重新生成 + assertNull(memoryMapper.selectById(memory1.getId())); + assertNull(memoryMapper.selectById(memory2.getId())); + assertNotNull(memoryMapper.selectById(memory3.getId())); + assertNotNull(insertMemory(1L, 100L).getId()); + } + + @Test + public void testDeleteMemoryListByIds_notExists() { + // mock 数据 + Ai1MemoryDO memory = insertMemory(1L, 100L); + + // 调用，并断言异常 + assertServiceException(() -> memoryService.deleteMemoryListByIds(Arrays.asList(memory.getId(), memory.getId() + 1)), + MEMORY_NOT_EXISTS); + // 断言 + assertNotNull(memoryMapper.selectById(memory.getId())); + } + + @Test + public void testDeleteMemoryListByAgentIds() { + // mock 数据 + Ai1MemoryDO memory1 = insertMemory(1L, 100L); + Ai1MemoryDO memory2 = insertMemory(2L, 100L); + Ai1MemoryDO memory3 = insertMemory(3L, 100L); + + // 调用 + memoryService.deleteMemoryListByAgentIds(Arrays.asList(1L, 2L)); + memoryService.deleteMemoryListByAgentIds(Collections.emptyList()); + // 断言 + assertNull(memoryMapper.selectById(memory1.getId())); + assertNull(memoryMapper.selectById(memory2.getId())); + assertNotNull(memoryMapper.selectById(memory3.getId())); + } + + @Test + public void testGetMemoryPage() { + // mock 数据 + Ai1MemoryDO memory1 = insertMemory(1L, 100L); + Ai1MemoryDO memory2 = insertMemory(1L, 200L); + Ai1MemoryDO memory3 = insertMemory(2L, 100L); + memoryService.deleteMemory(insertMemory(1L, 300L).getId()); + + // 调用：不传条件，按编号倒序，不含已删除的记忆 + PageResult pageResult = memoryService.getMemoryPage(new Ai1MemoryPageReqVO()); + // 断言 + assertEquals(List.of(memory3.getId(), memory2.getId(), memory1.getId()), + pageResult.getList().stream().map(Ai1MemoryDO::getId).toList()); + + // 调用：按 Agent、用户筛选 + pageResult = memoryService.getMemoryPage(new Ai1MemoryPageReqVO().setAgentId(1L).setUserId(100L)); + // 断言 + assertEquals(List.of(memory1.getId()), pageResult.getList().stream().map(Ai1MemoryDO::getId).toList()); + } + + @Test + public void testGetMemoryByAgentIdAndUserId() { + // mock 数据 + Ai1MemoryDO memory = insertMemory(1L, 100L); + insertMemory(1L, 200L); + + // 调用，并断言 + assertEquals(memory.getId(), memoryService.getMemoryByAgentIdAndUserId(1L, 100L).getId()); + assertNull(memoryService.getMemoryByAgentIdAndUserId(2L, 100L)); + } + + @Test + public void testGetMemoryByAgentIdAndUserId_duplicate() { + // mock 数据：极端并发下同一用户产生了重复记录 + insertMemory(1L, 100L); + insertMemory(1L, 100L); + + // 调用，并断言：不报错，返回其中一条 + assertNotNull(memoryService.getMemoryByAgentIdAndUserId(1L, 100L)); + } + + @Test + public void testCreateMemoryIfAgentExists() { + // mock 方法：Agent 1 存在，Agent 2 已删除 + when(agentService.getAgent(1L)).thenReturn(Ai1AgentDO.builder().id(1L).name("客服") + .status(CommonStatusEnum.ENABLE.getStatus()).build()); + + // 调用，并断言：Agent 已删除时不创建 + assertFalse(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(2L) + .userId(100L).content("迟到的归档结果").lastMessageId(20L).build())); + assertNull(memoryMapper.selectByAgentIdAndUserId(2L, 100L)); + // 调用，并断言：Agent 存在时创建 + assertTrue(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L) + .userId(100L).content("用户偏好").lastMessageId(20L).build())); + assertEquals("用户偏好", memoryMapper.selectByAgentIdAndUserId(1L, 100L).getContent()); + // 调用，并断言：已有未删除的记忆时不重复创建 + assertFalse(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L) + .userId(100L).content("重复归档").lastMessageId(30L).build())); + assertEquals(1, memoryMapper.selectCount(Ai1MemoryDO::getAgentId, 1L)); + } + + @Test + public void testCreateMemoryIfAgentExists_afterDelete() { + // mock 数据：已删除的历史记忆 + when(agentService.getAgent(1L)).thenReturn(Ai1AgentDO.builder().id(1L).name("客服") + .status(CommonStatusEnum.ENABLE.getStatus()).build()); + Ai1MemoryDO old = insertMemory(1L, 100L); + memoryService.deleteMemory(old.getId()); + + // 调用：删除后按历史消息重新生成 + assertTrue(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L) + .userId(100L).content("重建记忆").lastMessageId(5L).build())); + // 断言：生成新的记录与水位线，历史记录不影响按用户查询 + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(1L, 100L); + assertNotEquals(old.getId(), memory.getId()); + assertEquals("重建记忆", memory.getContent()); + assertEquals(5L, memory.getLastMessageId()); + // 断言：旧快照的归档写入不能复活已删除的记录 + assertFalse(memoryService.updateMemoryIfUnchanged(old, "迟到的归档结果", 30L)); + assertNull(memoryService.getMemory(old.getId())); + } + + @Test + public void testUpdateMemoryIfUnchanged() { + // mock 数据 + Ai1MemoryDO original = insertMemory(1L, 100L); + + // 调用，并断言：首次写入成功，按旧快照再次写入失败 + assertTrue(memoryService.updateMemoryIfUnchanged(original, "新记忆", 20L)); + assertFalse(memoryService.updateMemoryIfUnchanged(original, "过期结果", 30L)); + Ai1MemoryDO memory = memoryMapper.selectById(original.getId()); + assertEquals("新记忆", memory.getContent()); + assertEquals(20L, memory.getLastMessageId()); + } + + @Test + public void testUpdateMemoryIfUnchanged_edited() { + // mock 数据：归档读取后，人工只在末尾加了空格，水位线不变 + Ai1MemoryDO original = insertMemory(1L, 100L); + memoryMapper.updateById(new Ai1MemoryDO().setId(original.getId()).setContent("## 用户画像\n- 旧内容 ")); + + // 调用，并断言：放弃写入，保留人工修改 + assertFalse(memoryService.updateMemoryIfUnchanged(original, "归档结果", 20L)); + Ai1MemoryDO memory = memoryMapper.selectById(original.getId()); + assertEquals("## 用户画像\n- 旧内容 ", memory.getContent()); + assertEquals(10L, memory.getLastMessageId()); + } + + // ========== 测试数据 ========== + + private Ai1MemoryDO insertMemory(Long agentId, Long userId) { + Ai1MemoryDO memory = Ai1MemoryDO.builder().agentId(agentId).userId(userId) + .content("## 用户画像\n- 旧内容").lastMessageId(10L).build(); + memoryMapper.insert(memory); + return memory; + } + +} diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java index 63f5d7b..4605f25 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java @@ -1,11 +1,14 @@ package cn.iocoder.yudao.module.ai1.service.session; +import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder; import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO; import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMessageMapper; +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum; import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum; import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService; @@ -17,6 +20,9 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; +import org.mockito.InOrder; +import org.redisson.api.RLock; +import org.redisson.api.RedissonClient; import org.redisson.client.RedisConnectionException; import org.redisson.client.RedisTimeoutException; import org.springframework.context.annotation.Import; @@ -24,9 +30,11 @@ import org.springframework.test.context.bean.override.mockito.MockitoBean; import org.springframework.test.context.bean.override.mockito.MockitoSpyBean; import org.springframework.web.servlet.mvc.method.annotation.SseEmitter; +import java.util.ArrayList; import java.util.List; import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception; +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException; import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.SESSION_NOT_EXISTS; import static org.junit.jupiter.api.Assertions.*; import static org.mockito.ArgumentMatchers.*; @@ -37,7 +45,7 @@ import static org.mockito.Mockito.*; * * @author 芋道源码 */ -@Import(Ai1SessionMessageServiceImpl.class) +@Import({Ai1SessionMessageServiceImpl.class, Ai1MemoryLockRedisDAO.class}) public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { private static final Long TENANT_ID = 1L; @@ -55,11 +63,17 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { private Ai1AgentService agentService; @MockitoBean private Ai1SessionStreamTool sessionStreamTool; + @MockitoBean + private RedissonClient redissonClient; + + private final RLock snapshotLock = mock(RLock.class); @BeforeEach public void before() { // 设置租户上下文 TenantContextHolder.setTenantId(TENANT_ID); + // mock 记忆快照锁 + when(redissonClient.getLock(anyString())).thenReturn(snapshotLock); } @AfterEach @@ -100,6 +114,71 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { } @Test + public void testSendSessionMessageStream_snapshotLockUntilCommitted() { + // mock sessionService 的方法 + when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO() + .setId(10L).setAgentId(1L).setUserId(USER_ID).setTitle(Ai1SessionDO.TITLE_DEFAULT)); + // mock 释放快照锁时，本轮消息已提交 + List committedCounts = new ArrayList<>(); + doAnswer(invocation -> { + committedCounts.add(sessionMessageMapper.selectListBySessionId(10L).size()); + return null; + }).when(snapshotLock).unlock(); + + // 调用 + sessionMessageService.sendSessionMessageStream(USER_ID, + new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好")); + // 断言：按 Agent、用户加锁，落库后才释放，再投递任务 + verify(redissonClient).getLock("ai1:memory:snapshot-lock:1:" + USER_ID); + assertEquals(List.of(2), committedCounts); + InOrder inOrder = inOrder(snapshotLock, sessionStreamTool); + inOrder.verify(snapshotLock).lock(); + inOrder.verify(snapshotLock).unlock(); + inOrder.verify(sessionStreamTool).submit(anyLong(), eq(1L), eq(10L), eq("你好")); + } + + @Test + public void testStopSessionMessageMy_success() { + // 调用 + sessionMessageService.stopSessionMessageMy(USER_ID, 10L); + // 断言 + InOrder inOrder = inOrder(sessionService, sessionStreamTool); + inOrder.verify(sessionService).validateSessionMy(USER_ID, 10L); + inOrder.verify(sessionStreamTool).cancel(10L); + } + + @Test + public void testStopSessionMessageMy_notOwner() { + // mock sessionService 的方法 + when(sessionService.validateSessionMy(USER_ID, 10L)).thenThrow(exception(SESSION_NOT_EXISTS)); + + // 调用，并断言异常 + assertServiceException(() -> sessionMessageService.stopSessionMessageMy(USER_ID, 10L), SESSION_NOT_EXISTS); + // 断言 + verify(sessionStreamTool, never()).cancel(anyLong()); + } + + @Test + public void testGetSessionMessagePage() { + // mock 数据 + Ai1SessionMessageDO message1 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole()); + Ai1SessionMessageDO message2 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole()); + Ai1SessionMessageDO message3 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole()); + insertSessionMessage(11L, Ai1SessionMessageRoleEnum.USER.getRole()); + // 准备参数 + Ai1SessionMessagePageReqVO reqVO = new Ai1SessionMessagePageReqVO().setSessionId(10L); + reqVO.setPageSize(2); + + // 调用 + PageResult pageResult = sessionMessageService.getSessionMessagePage(reqVO); + // 断言：按编号倒序 + assertEquals(3, pageResult.getTotal()); + assertEquals(List.of(message3.getId(), message2.getId()), + pageResult.getList().stream().map(Ai1SessionMessageDO::getId).toList()); + assertNotEquals(message1.getId(), pageResult.getList().get(1).getId()); + } + + @Test public void testSendSessionMessageStream_submitFailed() { // mock sessionService 和 sessionStreamTool 的方法 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO() @@ -122,10 +201,11 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { @ParameterizedTest @ValueSource(ints = {1, 2}) - public void testSendSessionMessageStream_submitTimeoutKeepsWorkerResult(int status) { + public void testSendSessionMessageStream_submitTimeoutAfterWorkerCompleted(int status) { // mock sessionService 和 sessionStreamTool 的方法 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO() .setId(10L).setAgentId(1L).setTitle(Ai1SessionDO.TITLE_DEFAULT)); + // mock 投递实际成功，worker 在投递超时前已回填终态 doAnswer(invocation -> { Long messageId = invocation.getArgument(0); sessionMessageMapper.updateById(new Ai1SessionMessageDO().setId(messageId) @@ -136,7 +216,7 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { // 调用 sessionMessageService.sendSessionMessageStream(USER_ID, new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好")); - // 断言已有终态和内容不被补偿覆盖 + // 断言：补偿只更新生成中的消息，已有终态和内容不被覆盖 Ai1SessionMessageDO assistant = sessionMessageMapper.selectListBySessionId(10L).get(1); assertEquals(status, assistant.getStatus()); assertEquals("worker 已回填", assistant.getContent()); @@ -145,28 +225,6 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { } @Test - public void testSendSessionMessageStream_workerCanCompleteAfterSubmitTimeout() { - // mock sessionService 和 sessionStreamTool 的方法 - when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO() - .setId(10L).setAgentId(1L).setTitle(Ai1SessionDO.TITLE_DEFAULT)); - doThrow(new RedisTimeoutException("投递响应丢失")).when(sessionStreamTool) - .submit(anyLong(), eq(1L), eq(10L), eq("你好")); - - // 调用 - sessionMessageService.sendSessionMessageStream(USER_ID, - new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好")); - Ai1SessionMessageDO assistant = sessionMessageMapper.selectListBySessionId(10L).get(1); - assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), assistant.getStatus()); - // 调用 worker 回填，模拟 Redis 实际写入后响应丢失 - sessionMessageService.updateSessionMessage(new Ai1SessionMessageDO().setId(assistant.getId()) - .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus()).setContent("稍后生成成功")); - // 断言后续完成结果仍可落库 - Ai1SessionMessageDO result = sessionMessageMapper.selectById(assistant.getId()); - assertEquals(Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), result.getStatus()); - assertEquals("稍后生成成功", result.getContent()); - } - - @Test public void testSendSessionMessageStream_openFailedKeepsGenerating() { // mock sessionService 和 sessionStreamTool 的方法 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO() @@ -322,11 +380,71 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest { assertNotEquals(message1.getId(), list.get(1).getId()); } + @Test + public void testGetSessionMessageListBySessionIdsAndStatusAndIdBetween_safePrefix() { + // mock 数据：同一用户两个会话，第一个会话的助手消息仍在生成 + String user = Ai1SessionMessageRoleEnum.USER.getRole(); + String assistant = Ai1SessionMessageRoleEnum.ASSISTANT.getRole(); + Ai1SessionMessageDO message1 = insertSessionMessage(10L, user); + Ai1SessionMessageDO message2 = insertSessionMessage(10L, assistant, Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + Ai1SessionMessageDO message3 = insertSessionMessage(11L, user); + insertSessionMessage(12L, user); + insertSessionMessage(11L, assistant, Ai1SessionMessageStatusEnum.FAILED.getStatus()); + List sessionIds = List.of(10L, 11L); + + // 调用：只读取第一条生成中消息之前的已完成消息 + Ai1SessionMessageDO generating = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(sessionIds, + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), 0L); + List first = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 0L, generating.getId(), 100); + // 断言 + assertEquals(message2.getId(), generating.getId()); + assertEquals(List.of(message1.getId()), first.stream().map(Ai1SessionMessageDO::getId).toList()); + + // mock 生成结束 + sessionMessageMapper.updateById(new Ai1SessionMessageDO().setId(message2.getId()) + .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus())); + // 调用：从水位线之后读取，并按条数截断 + assertNull(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(sessionIds, + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), message1.getId())); + List limited = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message1.getId(), null, 1); + List all = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds, + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message1.getId(), null, 100); + // 断言 + assertEquals(List.of(message2.getId()), limited.stream().map(Ai1SessionMessageDO::getId).toList()); + assertEquals(List.of(message2.getId(), message3.getId()), all.stream().map(Ai1SessionMessageDO::getId).toList()); + } + + @Test + public void testGetGeneratingAssistantMessageIdList() { + // mock 数据：只有本会话生成中的助手消息命中 + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole()); + Ai1SessionMessageDO message1 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), + Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole()); + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.FAILED.getStatus()); + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + Ai1SessionMessageDO message2 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), + Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + insertSessionMessage(11L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus()); + + // 调用 + List ids = sessionMessageService.getGeneratingAssistantMessageIdList(10L); + // 断言：只返回本会话生成中的助手消息编号，按编号升序 + assertEquals(List.of(message1.getId(), message2.getId()), ids); + assertTrue(sessionMessageService.getGeneratingAssistantMessageIdList(12L).isEmpty()); + } + // ========== 测试数据 ========== private Ai1SessionMessageDO insertSessionMessage(Long sessionId, String role) { + return insertSessionMessage(sessionId, role, Ai1SessionMessageStatusEnum.SUCCESS.getStatus()); + } + + private Ai1SessionMessageDO insertSessionMessage(Long sessionId, String role, Integer status) { Ai1SessionMessageDO message = new Ai1SessionMessageDO().setSessionId(sessionId).setRole(role).setContent("内容") - .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus()); + .setStatus(status); sessionMessageMapper.insert(message); return message; } diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java index d8201bc..2fd899e 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java @@ -1,7 +1,9 @@ package cn.iocoder.yudao.module.ai1.service.session; +import cn.iocoder.yudao.framework.common.pojo.PageResult; import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionCreateMyReqVO; +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionPageReqVO; import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionUpdateMyReqVO; import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO; import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMapper; @@ -104,4 +106,44 @@ public class Ai1SessionServiceImplTest extends BaseDbUnitTest { assertEquals(Arrays.asList(id2, id1), list.stream().map(Ai1SessionDO::getId).toList()); } + @Test + public void testGetUserIdListByAgentId() { + // mock 数据 + sessionService.createSessionMy(200L, new Ai1SessionCreateMyReqVO().setAgentId(1L)); + sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L)); + sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L)); + sessionService.createSessionMy(300L, new Ai1SessionCreateMyReqVO().setAgentId(2L)); + + // 调用，并断言：去重并按用户编号升序 + assertEquals(Arrays.asList(100L, 200L), sessionService.getUserIdListByAgentId(1L)); + } + + @Test + public void testGetSession() { + // mock 数据 + Long id = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L)); + + // 调用，并断言 + assertEquals(100L, sessionService.getSession(id).getUserId()); + assertNull(sessionService.getSession(id + 1)); + } + + @Test + public void testGetSessionPage_withoutAgentId() { + // mock 数据 + Long id1 = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L)); + Long id2 = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(2L)); + Long id3 = sessionService.createSessionMy(200L, new Ai1SessionCreateMyReqVO().setAgentId(2L)); + + // 调用：不传 Agent 时跨 Agent 查询 + PageResult pageResult = sessionService.getSessionPage(new Ai1SessionPageReqVO()); + // 断言 + assertEquals(Arrays.asList(id3, id2, id1), pageResult.getList().stream().map(Ai1SessionDO::getId).toList()); + + // 调用：按 Agent、用户筛选 + pageResult = sessionService.getSessionPage(new Ai1SessionPageReqVO().setAgentId(2L).setUserId(100L)); + // 断言 + assertEquals(Collections.singletonList(id2), pageResult.getList().stream().map(Ai1SessionDO::getId).toList()); + } + } diff --git a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/util/Ai1UtilsTest.java b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/util/Ai1UtilsTest.java index 514ba0f..1eb0281 100644 --- a/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/util/Ai1UtilsTest.java +++ b/yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/util/Ai1UtilsTest.java @@ -17,6 +17,8 @@ import static org.mockito.Mockito.mockStatic; public class Ai1UtilsTest { + private static final List MEMORY_SECTIONS = List.of("用户画像", "偏好习惯", "关键事实", "目标与意图", "近期进展"); + @Test public void testNormalizeHeaders_orderAndValues() { // 准备参数 @@ -116,4 +118,55 @@ public class Ai1UtilsTest { assertFalse(ex.getMessage().contains("literal-secret")); } } + + @Test + public void testNormalizeMemoryMarkdown_sections() { + // 调用，并断言：标题与内容压成一行时拆开，不同维度之间空一行 + assertEquals("## 用户画像\n- 后端工程师\n\n## 偏好习惯\n- 简洁回答", + Ai1Utils.normalizeMemoryMarkdown("##用户画像-后端工程师##偏好习惯-简洁回答", MEMORY_SECTIONS)); + // 调用，并断言：统一换行，合并多余空行 + assertEquals("## 用户画像\n- 程序员\n\n## 关键事实\n- 住在北京", + Ai1Utils.normalizeMemoryMarkdown(" ## 用户画像\r\n-程序员\r\n\r\n\r\n## 关键事实\r- 住在北京 ", MEMORY_SECTIONS)); + } + + @Test + public void testNormalizeMemoryMarkdown_keepsFacts() { + // 准备参数：C#、#123、URL 锚点、中文破折号都是用户事实，不能被拆开 + String content = "## 关键事实\n- 主力语言 C#，工单#123，文档 https://a.cn/doc#part\n- 常驻北京—上海，周期 3–5 天"; + + // 调用，并断言 + assertEquals(content, Ai1Utils.normalizeMemoryMarkdown(content, MEMORY_SECTIONS)); + // 调用，并断言：非固定维度的行内标题不拆行 + assertEquals("## 关键事实\n- 标签##其他", Ai1Utils.normalizeMemoryMarkdown("## 关键事实\n- 标签##其他", MEMORY_SECTIONS)); + } + + @Test + public void testNormalizeMemoryMarkdown_inlineList() { + // 调用，并断言：非 ASCII 字符后紧跟 ASCII「-」，拆成列表项 + assertEquals("## 偏好习惯\n- 口吻正式\n- 开场先问候", + Ai1Utils.normalizeMemoryMarkdown("## 偏好习惯\n- 口吻正式-开场先问候", MEMORY_SECTIONS)); + // 调用，并断言：ASCII 字符之间的「-」不拆 + assertEquals("## 关键事实\n- 使用 spring-boot", Ai1Utils.normalizeMemoryMarkdown("## 关键事实\n- 使用 spring-boot", MEMORY_SECTIONS)); + } + + @Test + public void testNormalizeMemoryMarkdown_blank() { + // 调用，并断言：为空时原样返回 + assertNull(Ai1Utils.normalizeMemoryMarkdown(null, MEMORY_SECTIONS)); + assertEquals(" ", Ai1Utils.normalizeMemoryMarkdown(" ", MEMORY_SECTIONS)); + } + + @Test + public void testLimitLength() { + // 调用，并断言：先去除首尾空白，未超长时不截断 + assertEquals("用户喜欢", Ai1Utils.limitLength(" 用户喜欢 \n", 4)); + // 调用，并断言：下标 4、5 的 char 是同一个 emoji，截断不拆开代理对 + assertEquals("用户喜欢", Ai1Utils.limitLength("用户喜欢😀表情", 5)); + assertEquals("用户喜欢😀", Ai1Utils.limitLength("用户喜欢😀表情", 6)); + assertEquals("用户喜欢😀表", Ai1Utils.limitLength("用户喜欢😀表情", 7)); + // 调用，并断言：为空时原样返回 + assertNull(Ai1Utils.limitLength(null, 5)); + assertEquals(" ", Ai1Utils.limitLength(" ", 5)); + } + } diff --git a/yudao-module-ai1-springai/src/test/resources/sql/clean.sql b/yudao-module-ai1-springai/src/test/resources/sql/clean.sql index a48e1e2..f774b1c 100644 --- a/yudao-module-ai1-springai/src/test/resources/sql/clean.sql +++ b/yudao-module-ai1-springai/src/test/resources/sql/clean.sql @@ -8,3 +8,4 @@ DELETE FROM "ai1_knowledge_document"; DELETE FROM "ai1_agent"; DELETE FROM "ai1_session"; DELETE FROM "ai1_session_message"; +DELETE FROM "ai1_memory"; diff --git a/yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql b/yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql index c4ef951..ff42361 100644 --- a/yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql +++ b/yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql @@ -152,6 +152,7 @@ CREATE TABLE IF NOT EXISTS "ai1_agent" ( "knowledge_base_ids" varchar(500) DEFAULT NULL, "mcp_ids" varchar(500) DEFAULT NULL, "skill_ids" varchar(500) DEFAULT NULL, + "memory_enabled" bit NOT NULL DEFAULT FALSE, "status" tinyint NOT NULL DEFAULT 1, "creator" varchar(64) DEFAULT '', "create_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, @@ -197,3 +198,21 @@ CREATE TABLE IF NOT EXISTS "ai1_session_message" ( "tenant_id" bigint NOT NULL DEFAULT 0, PRIMARY KEY ("id") ); + +-- ---------------------------- +-- AI1 长期记忆 +-- ---------------------------- +CREATE TABLE IF NOT EXISTS "ai1_memory" ( + "id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY, + "agent_id" bigint NOT NULL, + "user_id" bigint NOT NULL, + "content" mediumtext NOT NULL, + "last_message_id" bigint NOT NULL DEFAULT 0, + "creator" varchar(64) DEFAULT '', + "create_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updater" varchar(64) DEFAULT '', + "update_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + "deleted" bit NOT NULL DEFAULT FALSE, + "tenant_id" bigint NOT NULL DEFAULT 0, + PRIMARY KEY ("id") +); -- libgit2 0.21.4