Commit edffc81991441eef863618d779ed8358c13f90c2

Authored by YunaiV
1 parent 17917314

feat(ai1): 增加长期记忆并完善会话管理与生成控制

- 增加长期记忆管理、归档与租户 Job,支持记忆筛选和人工维护
- 支持跨 Agent 会话查询、消息分页和按状态筛选 Agent
- 完善停止生成、断线续传、Redis 恢复及技能进程回收
- 按项目风格整理相关实现与测试
Showing 52 changed files with 3885 additions and 103 deletions
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/agent/Ai1AgentController.java
1 1 package cn.iocoder.yudao.module.ai1.controller.admin.agent;
2 2  
3 3 import cn.hutool.core.collection.CollUtil;
4   -import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
5 4 import cn.iocoder.yudao.framework.common.pojo.CommonResult;
6 5 import cn.iocoder.yudao.framework.common.pojo.PageResult;
7 6 import cn.iocoder.yudao.framework.common.util.collection.MapUtils;
... ... @@ -97,9 +96,11 @@ public class Ai1AgentController {
97 96 }
98 97  
99 98 @GetMapping("/simple-list")
100   - @Operation(summary = "获得已开启的 Agent 精简列表", description = "用于会话页选择 Agent,仅返回已开启的 Agent")
101   - public CommonResult<List<Ai1AgentRespVO>> getAgentSimpleList() {
102   - List<Ai1AgentDO> list = agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus());
  99 + @Operation(summary = "获得 Agent 精简列表", description = "按状态筛选,不传时返回全部状态;会话页传 0 仅选择已开启的 Agent")
  100 + @Parameter(name = "status", description = "状态,参见 CommonStatusEnum 枚举:0 开启,1 关闭;不传时返回全部状态", example = "0")
  101 + public CommonResult<List<Ai1AgentRespVO>> getAgentSimpleList(
  102 + @RequestParam(value = "status", required = false) Integer status) {
  103 + List<Ai1AgentDO> list = agentService.getAgentListByStatus(status);
103 104 Map<Long, Ai1ModelDO> modelMap = modelService.getModelMap(convertSet(list, Ai1AgentDO::getModelId));
104 105 return success(convertList(list, agent -> {
105 106 Ai1AgentRespVO respVO = new Ai1AgentRespVO().setId(agent.getId()).setName(agent.getName())
... ...
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 {
43 43 @Schema(description = "SKILL 编号集合", example = "[1]")
44 44 private List<Long> skillIds;
45 45  
  46 + @Schema(description = "是否开启长期记忆", requiredMode = Schema.RequiredMode.REQUIRED, example = "false")
  47 + private Boolean memoryEnabled;
  48 +
46 49 @Schema(description = "状态", requiredMode = Schema.RequiredMode.REQUIRED, example = "0")
47 50 private Integer status;
48 51  
... ...
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 {
44 44 @Schema(description = "SKILL 编号集合", example = "[1]")
45 45 private List<Long> skillIds;
46 46  
  47 + @Schema(description = "是否开启长期记忆", example = "false")
  48 + private Boolean memoryEnabled;
  49 +
47 50 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/Ai1MemoryController.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.controller.admin.memory;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  5 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  6 +import cn.iocoder.yudao.framework.common.util.collection.MapUtils;
  7 +import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  8 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO;
  9 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryRespVO;
  10 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO;
  11 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  12 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  13 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
  14 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
  15 +import cn.iocoder.yudao.module.system.api.user.AdminUserApi;
  16 +import cn.iocoder.yudao.module.system.api.user.dto.AdminUserRespDTO;
  17 +import io.swagger.v3.oas.annotations.Operation;
  18 +import io.swagger.v3.oas.annotations.Parameter;
  19 +import io.swagger.v3.oas.annotations.tags.Tag;
  20 +import jakarta.annotation.Resource;
  21 +import jakarta.validation.Valid;
  22 +import org.springframework.security.access.prepost.PreAuthorize;
  23 +import org.springframework.validation.annotation.Validated;
  24 +import org.springframework.web.bind.annotation.*;
  25 +
  26 +import java.util.Collections;
  27 +import java.util.List;
  28 +import java.util.Map;
  29 +
  30 +import static cn.iocoder.yudao.framework.common.pojo.CommonResult.success;
  31 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertSet;
  32 +
  33 +@Tag(name = "管理后台 - AI1 长期记忆")
  34 +@RestController
  35 +@RequestMapping("/ai1/memory")
  36 +@Validated
  37 +public class Ai1MemoryController {
  38 +
  39 + @Resource
  40 + private Ai1MemoryService memoryService;
  41 + @Resource
  42 + private Ai1AgentService agentService;
  43 +
  44 + @Resource
  45 + private AdminUserApi adminUserApi;
  46 +
  47 + @PutMapping("/update")
  48 + @Operation(summary = "更新长期记忆")
  49 + @PreAuthorize("@ss.hasPermission('ai1:memory:update')")
  50 + public CommonResult<Boolean> updateMemory(@Valid @RequestBody Ai1MemoryUpdateReqVO updateReqVO) {
  51 + memoryService.updateMemory(updateReqVO);
  52 + return success(true);
  53 + }
  54 +
  55 + @DeleteMapping("/delete")
  56 + @Operation(summary = "删除长期记忆", description = "之后根据历史对话重新生成")
  57 + @Parameter(name = "id", description = "编号", required = true)
  58 + @PreAuthorize("@ss.hasPermission('ai1:memory:delete')")
  59 + public CommonResult<Boolean> deleteMemory(@RequestParam("id") Long id) {
  60 + memoryService.deleteMemory(id);
  61 + return success(true);
  62 + }
  63 +
  64 + @DeleteMapping("/delete-list")
  65 + @Operation(summary = "批量删除长期记忆", description = "之后根据历史对话重新生成")
  66 + @Parameter(name = "ids", description = "编号列表", required = true)
  67 + @PreAuthorize("@ss.hasPermission('ai1:memory:delete')")
  68 + public CommonResult<Boolean> deleteMemoryList(@RequestParam("ids") List<Long> ids) {
  69 + memoryService.deleteMemoryListByIds(ids);
  70 + return success(true);
  71 + }
  72 +
  73 + @GetMapping("/get")
  74 + @Operation(summary = "获得长期记忆")
  75 + @Parameter(name = "id", description = "编号", required = true, example = "1024")
  76 + @PreAuthorize("@ss.hasPermission('ai1:memory:query')")
  77 + public CommonResult<Ai1MemoryRespVO> getMemory(@RequestParam("id") Long id) {
  78 + Ai1MemoryDO memory = memoryService.getMemory(id);
  79 + return success(buildMemoryRespVO(memory));
  80 + }
  81 +
  82 + @GetMapping("/page")
  83 + @Operation(summary = "获得长期记忆分页")
  84 + @PreAuthorize("@ss.hasPermission('ai1:memory:query')")
  85 + public CommonResult<PageResult<Ai1MemoryRespVO>> getMemoryPage(@Valid Ai1MemoryPageReqVO pageReqVO) {
  86 + PageResult<Ai1MemoryDO> pageResult = memoryService.getMemoryPage(pageReqVO);
  87 + return success(new PageResult<>(buildMemoryRespVOList(pageResult.getList()), pageResult.getTotal()));
  88 + }
  89 +
  90 + // ==================== 拼接 VO ====================
  91 +
  92 + private Ai1MemoryRespVO buildMemoryRespVO(Ai1MemoryDO memory) {
  93 + if (memory == null) {
  94 + return null;
  95 + }
  96 + return CollUtil.getFirst(buildMemoryRespVOList(Collections.singletonList(memory)));
  97 + }
  98 +
  99 + private List<Ai1MemoryRespVO> buildMemoryRespVOList(List<Ai1MemoryDO> list) {
  100 + if (CollUtil.isEmpty(list)) {
  101 + return Collections.emptyList();
  102 + }
  103 + Map<Long, Ai1AgentDO> agentMap = agentService.getAgentMap(convertSet(list, Ai1MemoryDO::getAgentId));
  104 + Map<Long, AdminUserRespDTO> userMap = adminUserApi.getUserMap(convertSet(list, Ai1MemoryDO::getUserId));
  105 + return BeanUtils.toBean(list, Ai1MemoryRespVO.class, respVO -> {
  106 + MapUtils.findAndThen(agentMap, respVO.getAgentId(), agent -> respVO.setAgentName(agent.getName()));
  107 + MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname()));
  108 + });
  109 + }
  110 +
  111 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryPageReqVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.PageParam;
  4 +import io.swagger.v3.oas.annotations.media.Schema;
  5 +import lombok.Data;
  6 +import lombok.EqualsAndHashCode;
  7 +import lombok.ToString;
  8 +
  9 +@Schema(description = "管理后台 - AI1 长期记忆分页 Request VO")
  10 +@Data
  11 +@EqualsAndHashCode(callSuper = true)
  12 +@ToString(callSuper = true)
  13 +public class Ai1MemoryPageReqVO extends PageParam {
  14 +
  15 + @Schema(description = "Agent 编号", example = "1")
  16 + private Long agentId;
  17 +
  18 + @Schema(description = "用户编号", example = "1")
  19 + private Long userId;
  20 +
  21 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryRespVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo;
  2 +
  3 +import io.swagger.v3.oas.annotations.media.Schema;
  4 +import lombok.Data;
  5 +
  6 +import java.time.LocalDateTime;
  7 +
  8 +@Schema(description = "管理后台 - AI1 长期记忆 Response VO")
  9 +@Data
  10 +public class Ai1MemoryRespVO {
  11 +
  12 + @Schema(description = "编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024")
  13 + private Long id;
  14 +
  15 + @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
  16 + private Long agentId;
  17 +
  18 + @Schema(description = "Agent 名称", example = "客服助手")
  19 + private String agentName;
  20 +
  21 + @Schema(description = "用户编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
  22 + private Long userId;
  23 +
  24 + @Schema(description = "用户昵称", example = "芋道源码")
  25 + private String userNickname;
  26 +
  27 + @Schema(description = "记忆正文", requiredMode = Schema.RequiredMode.REQUIRED, example = "## 用户画像\n- 后端工程师")
  28 + private String content;
  29 +
  30 + @Schema(description = "已归档的最后一条消息编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "2048")
  31 + private Long lastMessageId;
  32 +
  33 + @Schema(description = "创建时间", requiredMode = Schema.RequiredMode.REQUIRED)
  34 + private LocalDateTime createTime;
  35 +
  36 + @Schema(description = "更新时间", requiredMode = Schema.RequiredMode.REQUIRED)
  37 + private LocalDateTime updateTime;
  38 +
  39 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/memory/vo/Ai1MemoryUpdateReqVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.controller.admin.memory.vo;
  2 +
  3 +import io.swagger.v3.oas.annotations.media.Schema;
  4 +import jakarta.validation.constraints.NotBlank;
  5 +import jakarta.validation.constraints.NotNull;
  6 +import jakarta.validation.constraints.Size;
  7 +import lombok.Data;
  8 +
  9 +@Schema(description = "管理后台 - AI1 长期记忆修改 Request VO")
  10 +@Data
  11 +public class Ai1MemoryUpdateReqVO {
  12 +
  13 + @Schema(description = "编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024")
  14 + @NotNull(message = "编号不能为空")
  15 + private Long id;
  16 +
  17 + @Schema(description = "记忆正文", requiredMode = Schema.RequiredMode.REQUIRED, example = "## 用户画像\n- 后端工程师")
  18 + @NotBlank(message = "记忆正文不能为空")
  19 + @Size(max = 65535, message = "记忆正文长度不能超过 65535 个字符")
  20 + private String content;
  21 +
  22 +}
... ...
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
8 8 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionPageReqVO;
9 9 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionRespVO;
10 10 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionUpdateMyReqVO;
  11 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
11 12 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
  13 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
12 14 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
13 15 import cn.iocoder.yudao.module.system.api.user.AdminUserApi;
14 16 import cn.iocoder.yudao.module.system.api.user.dto.AdminUserRespDTO;
... ... @@ -36,6 +38,8 @@ public class Ai1SessionController {
36 38  
37 39 @Resource
38 40 private Ai1SessionService sessionService;
  41 + @Resource
  42 + private Ai1AgentService agentService;
39 43  
40 44 @Resource
41 45 private AdminUserApi adminUserApi;
... ... @@ -78,13 +82,16 @@ public class Ai1SessionController {
78 82 // ==================== 管理 ====================
79 83  
80 84 @GetMapping("/page")
81   - @Operation(summary = "获得会话分页", description = "管理端查看 Agent 下全部用户的会话")
  85 + @Operation(summary = "获得会话分页", description = "管理端查看全部用户的会话,可按 Agent 筛选")
82 86 @PreAuthorize("@ss.hasPermission('ai1:session:query')")
83 87 public CommonResult<PageResult<Ai1SessionRespVO>> getSessionPage(@Valid Ai1SessionPageReqVO pageReqVO) {
84 88 PageResult<Ai1SessionDO> pageResult = sessionService.getSessionPage(pageReqVO);
  89 + Map<Long, Ai1AgentDO> agentMap = agentService.getAgentMap(convertSet(pageResult.getList(), Ai1SessionDO::getAgentId));
85 90 Map<Long, AdminUserRespDTO> userMap = adminUserApi.getUserMap(convertSet(pageResult.getList(), Ai1SessionDO::getUserId));
86   - return success(BeanUtils.toBean(pageResult, Ai1SessionRespVO.class, respVO ->
87   - MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname()))));
  91 + return success(BeanUtils.toBean(pageResult, Ai1SessionRespVO.class, respVO -> {
  92 + MapUtils.findAndThen(agentMap, respVO.getAgentId(), agent -> respVO.setAgentName(agent.getName()));
  93 + MapUtils.findAndThen(userMap, respVO.getUserId(), user -> respVO.setUserNickname(user.getNickname()));
  94 + }));
88 95 }
89 96  
90 97 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/Ai1SessionMessageController.java
1 1 package cn.iocoder.yudao.module.ai1.controller.admin.session;
2 2  
3 3 import cn.iocoder.yudao.framework.common.pojo.CommonResult;
  4 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
4 5 import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  6 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO;
5 7 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageRespVO;
6 8 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO;
7 9 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
... ... @@ -63,8 +65,25 @@ public class Ai1SessionMessageController {
63 65 return sessionMessageService.resumeSessionMessageStream(getLoginUserId(), messageId, lastEventId);
64 66 }
65 67  
  68 + @PostMapping("/stop-my")
  69 + @Operation(summary = "停止【我的】会话生成", description = "终止会话中生成中的助手消息,保留已生成的内容")
  70 + @Parameter(name = "sessionId", description = "会话编号", required = true, example = "1024")
  71 + @PreAuthorize("@ss.hasPermission('ai1:session:my')")
  72 + public CommonResult<Boolean> stopSessionMessageMy(@RequestParam("sessionId") Long sessionId) {
  73 + sessionMessageService.stopSessionMessageMy(getLoginUserId(), sessionId);
  74 + return success(true);
  75 + }
  76 +
66 77 // ==================== 管理 ====================
67 78  
  79 + @GetMapping("/page")
  80 + @Operation(summary = "获得会话消息分页", description = "按消息编号倒序")
  81 + @PreAuthorize("@ss.hasPermission('ai1:session:query')")
  82 + public CommonResult<PageResult<Ai1SessionMessageRespVO>> getSessionMessagePage(@Valid Ai1SessionMessagePageReqVO pageReqVO) {
  83 + return success(BeanUtils.toBean(sessionMessageService.getSessionMessagePage(pageReqVO),
  84 + Ai1SessionMessageRespVO.class));
  85 + }
  86 +
68 87 @GetMapping("/list")
69 88 @Operation(summary = "获得会话消息列表", description = "管理端查看任意用户的会话明细")
70 89 @Parameter(name = "sessionId", description = "会话编号", required = true, example = "1024")
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/controller/admin/session/vo/message/Ai1SessionMessagePageReqVO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.PageParam;
  4 +import io.swagger.v3.oas.annotations.media.Schema;
  5 +import jakarta.validation.constraints.NotNull;
  6 +import lombok.Data;
  7 +import lombok.EqualsAndHashCode;
  8 +import lombok.ToString;
  9 +
  10 +@Schema(description = "管理后台 - AI1 会话消息分页 Request VO")
  11 +@Data
  12 +@EqualsAndHashCode(callSuper = true)
  13 +@ToString(callSuper = true)
  14 +public class Ai1SessionMessagePageReqVO extends PageParam {
  15 +
  16 + @Schema(description = "会话编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1024")
  17 + @NotNull(message = "会话编号不能为空")
  18 + private Long sessionId;
  19 +
  20 +}
... ...
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;
2 2  
3 3 import cn.iocoder.yudao.framework.common.pojo.PageParam;
4 4 import io.swagger.v3.oas.annotations.media.Schema;
5   -import jakarta.validation.constraints.NotNull;
6 5 import lombok.Data;
7 6 import lombok.EqualsAndHashCode;
8 7 import lombok.ToString;
... ... @@ -13,8 +12,7 @@ import lombok.ToString;
13 12 @ToString(callSuper = true)
14 13 public class Ai1SessionPageReqVO extends PageParam {
15 14  
16   - @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
17   - @NotNull(message = "Agent 编号不能为空")
  15 + @Schema(description = "Agent 编号,为空时查询全部 Agent", example = "1")
18 16 private Long agentId;
19 17  
20 18 @Schema(description = "标题,模糊匹配", example = "三体")
... ...
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 {
15 15 @Schema(description = "Agent 编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
16 16 private Long agentId;
17 17  
  18 + @Schema(description = "Agent 名称", example = "客服助手")
  19 + private String agentName;
  20 +
18 21 @Schema(description = "用户编号", requiredMode = Schema.RequiredMode.REQUIRED, example = "1")
19 22 private Long userId;
20 23  
... ...
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 {
82 82 @TableField(typeHandler = LongListTypeHandler.class)
83 83 private List<Long> skillIds;
84 84 /**
  85 + * 是否开启长期记忆
  86 + *
  87 + * 开启后,对话时注入记忆,并由 Job 归档已完成的消息
  88 + */
  89 + private Boolean memoryEnabled;
  90 + /**
85 91 * 状态
86 92 *
87 93 * 枚举 {@link CommonStatusEnum};开启后才可在「Agent 会话」中使用
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/dataobject/memory/Ai1MemoryDO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.dal.dataobject.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.tenant.core.db.TenantBaseDO;
  4 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  5 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
  6 +import com.baomidou.mybatisplus.annotation.KeySequence;
  7 +import com.baomidou.mybatisplus.annotation.TableId;
  8 +import com.baomidou.mybatisplus.annotation.TableName;
  9 +import lombok.*;
  10 +
  11 +/**
  12 + * AI1 长期记忆 DO
  13 + *
  14 + * 每个用户在每个 Agent 下一条未删除的记忆,由归档 Job 根据已完成的消息增量生成;
  15 + * 删除为逻辑删除,保留历史记录便于追溯,之后的归档会根据历史消息生成新的记录
  16 + *
  17 + * @author 芋道源码
  18 + */
  19 +@TableName("ai1_memory")
  20 +@KeySequence("ai1_memory_seq") // 用于 Oracle、PostgreSQL、Kingbase、DB2、H2 数据库的主键自增。如果是 MySQL 等数据库,可不写。
  21 +@Data
  22 +@EqualsAndHashCode(callSuper = true)
  23 +@ToString(callSuper = true)
  24 +@Builder
  25 +@NoArgsConstructor
  26 +@AllArgsConstructor
  27 +public class Ai1MemoryDO extends TenantBaseDO {
  28 +
  29 + /**
  30 + * 编号
  31 + */
  32 + @TableId
  33 + private Long id;
  34 + /**
  35 + * Agent 编号
  36 + *
  37 + * 关联 {@link Ai1AgentDO#getId()}
  38 + */
  39 + private Long agentId;
  40 + /**
  41 + * 用户编号
  42 + *
  43 + * 关联 AdminUserDO 的 id 编号
  44 + */
  45 + private Long userId;
  46 + /**
  47 + * 记忆正文,Markdown 格式,按维度分节
  48 + */
  49 + private String content;
  50 + /**
  51 + * 已归档的最后一条消息编号,即增量归档的水位线
  52 + *
  53 + * 关联 {@link Ai1SessionMessageDO#getId()}
  54 + */
  55 + private Long lastMessageId;
  56 +
  57 +}
... ...
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&lt;Ai1AgentDO&gt; {
28 28  
29 29 default List<Ai1AgentDO> selectListByStatusOrderById(Integer status) {
30 30 return selectList(new LambdaQueryWrapperX<Ai1AgentDO>()
31   - .eq(Ai1AgentDO::getStatus, status)
  31 + .eqIfPresent(Ai1AgentDO::getStatus, status)
32 32 .orderByAsc(Ai1AgentDO::getId));
33 33 }
34 34  
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/memory/Ai1MemoryMapper.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.dal.mysql.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  4 +import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
  5 +import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
  6 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO;
  7 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  8 +import org.apache.ibatis.annotations.Mapper;
  9 +
  10 +/**
  11 + * AI1 长期记忆 Mapper
  12 + *
  13 + * @author 芋道源码
  14 + */
  15 +@Mapper
  16 +public interface Ai1MemoryMapper extends BaseMapperX<Ai1MemoryDO> {
  17 +
  18 + default PageResult<Ai1MemoryDO> selectPage(Ai1MemoryPageReqVO reqVO) {
  19 + return selectPage(reqVO, new LambdaQueryWrapperX<Ai1MemoryDO>()
  20 + .eqIfPresent(Ai1MemoryDO::getAgentId, reqVO.getAgentId())
  21 + .eqIfPresent(Ai1MemoryDO::getUserId, reqVO.getUserId())
  22 + .orderByDesc(Ai1MemoryDO::getId));
  23 + }
  24 +
  25 + /**
  26 + * 未加唯一索引,极端并发下可能存在重复记录,取第一条
  27 + */
  28 + default Ai1MemoryDO selectByAgentIdAndUserId(Long agentId, Long userId) {
  29 + return selectFirstOne(Ai1MemoryDO::getAgentId, agentId, Ai1MemoryDO::getUserId, userId);
  30 + }
  31 +
  32 +}
... ...
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;
10 10 import java.util.Collection;
11 11 import java.util.List;
12 12  
  13 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
  14 +
13 15 /**
14 16 * AI1 会话 Mapper
15 17 *
... ... @@ -20,7 +22,7 @@ public interface Ai1SessionMapper extends BaseMapperX&lt;Ai1SessionDO&gt; {
20 22  
21 23 default PageResult<Ai1SessionDO> selectPage(Ai1SessionPageReqVO reqVO) {
22 24 return selectPage(reqVO, new LambdaQueryWrapperX<Ai1SessionDO>()
23   - .eq(Ai1SessionDO::getAgentId, reqVO.getAgentId())
  25 + .eqIfPresent(Ai1SessionDO::getAgentId, reqVO.getAgentId())
24 26 .likeIfPresent(Ai1SessionDO::getTitle, reqVO.getTitle())
25 27 .eqIfPresent(Ai1SessionDO::getUserId, reqVO.getUserId())
26 28 .orderByDesc(Ai1SessionDO::getId));
... ... @@ -33,6 +35,14 @@ public interface Ai1SessionMapper extends BaseMapperX&lt;Ai1SessionDO&gt; {
33 35 .orderByDesc(Ai1SessionDO::getId));
34 36 }
35 37  
  38 + default List<Long> selectUserIdListByAgentId(Long agentId) {
  39 + return convertList(selectList(new LambdaQueryWrapperX<Ai1SessionDO>()
  40 + .eq(Ai1SessionDO::getAgentId, agentId)
  41 + .select(Ai1SessionDO::getUserId)
  42 + .groupBy(Ai1SessionDO::getUserId)
  43 + .orderByAsc(Ai1SessionDO::getUserId)), Ai1SessionDO::getUserId);
  44 + }
  45 +
36 46 default List<Ai1SessionDO> selectListByAgentIds(Collection<Long> agentIds) {
37 47 return selectList(new LambdaQueryWrapperX<Ai1SessionDO>().in(Ai1SessionDO::getAgentId, agentIds));
38 48 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/mysql/session/Ai1SessionMessageMapper.java
1 1 package cn.iocoder.yudao.module.ai1.dal.mysql.session;
2 2  
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
3 4 import cn.iocoder.yudao.framework.mybatis.core.mapper.BaseMapperX;
4 5 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
5 6 import cn.iocoder.yudao.framework.mybatis.core.query.QueryWrapperX;
  7 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO;
6 8 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
7 9 import org.apache.ibatis.annotations.Mapper;
8 10  
9 11 import java.util.Collection;
10 12 import java.util.List;
11 13  
  14 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
  15 +
12 16 /**
13 17 * AI1 会话消息 Mapper
14 18 *
... ... @@ -17,6 +21,12 @@ import java.util.List;
17 21 @Mapper
18 22 public interface Ai1SessionMessageMapper extends BaseMapperX<Ai1SessionMessageDO> {
19 23  
  24 + default PageResult<Ai1SessionMessageDO> selectPage(Ai1SessionMessagePageReqVO reqVO) {
  25 + return selectPage(reqVO, new LambdaQueryWrapperX<Ai1SessionMessageDO>()
  26 + .eq(Ai1SessionMessageDO::getSessionId, reqVO.getSessionId())
  27 + .orderByDesc(Ai1SessionMessageDO::getId));
  28 + }
  29 +
20 30 default List<Ai1SessionMessageDO> selectListBySessionId(Long sessionId) {
21 31 return selectList(new LambdaQueryWrapperX<Ai1SessionMessageDO>()
22 32 .eq(Ai1SessionMessageDO::getSessionId, sessionId)
... ... @@ -24,6 +34,18 @@ public interface Ai1SessionMessageMapper extends BaseMapperX&lt;Ai1SessionMessageDO
24 34 }
25 35  
26 36 /**
  37 + * 只查询编号,不读取正文、思考过程
  38 + */
  39 + default List<Long> selectIdListBySessionIdAndRoleAndStatus(Long sessionId, String role, Integer status) {
  40 + return convertList(selectList(new LambdaQueryWrapperX<Ai1SessionMessageDO>()
  41 + .select(Ai1SessionMessageDO::getId)
  42 + .eq(Ai1SessionMessageDO::getSessionId, sessionId)
  43 + .eq(Ai1SessionMessageDO::getRole, role)
  44 + .eq(Ai1SessionMessageDO::getStatus, status)
  45 + .orderByAsc(Ai1SessionMessageDO::getId)), Ai1SessionMessageDO::getId);
  46 + }
  47 +
  48 + /**
27 49 * 按编号倒序,取指定消息之前最近的 limit 条
28 50 */
29 51 default List<Ai1SessionMessageDO> selectListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit) {
... ... @@ -34,6 +56,32 @@ public interface Ai1SessionMessageMapper extends BaseMapperX&lt;Ai1SessionMessageDO
34 56 .orderByDesc("id"));
35 57 }
36 58  
  59 + /**
  60 + * 获得多个会话中,编号大于 id 的第一条指定状态的消息
  61 + */
  62 + default Ai1SessionMessageDO selectFirstBySessionIdsAndStatusAndIdGt(Collection<Long> sessionIds, Integer status, Long id) {
  63 + return selectOne(new QueryWrapperX<Ai1SessionMessageDO>()
  64 + .limitN(1)
  65 + .in("session_id", sessionIds)
  66 + .eq("status", status)
  67 + .gt("id", id)
  68 + .orderByAsc("id"));
  69 + }
  70 +
  71 + /**
  72 + * 获得多个会话中,编号在 (minId, maxId) 之间的指定状态消息,按编号升序取前 limit 条;maxId 为空时不限上界
  73 + */
  74 + default List<Ai1SessionMessageDO> selectListBySessionIdsAndStatusAndIdBetween(Collection<Long> sessionIds, Integer status,
  75 + Long minId, Long maxId, Integer limit) {
  76 + return selectList(new QueryWrapperX<Ai1SessionMessageDO>()
  77 + .limitN(limit)
  78 + .in("session_id", sessionIds)
  79 + .eq("status", status)
  80 + .gt("id", minId)
  81 + .lt(maxId != null, "id", maxId)
  82 + .orderByAsc("id"));
  83 + }
  84 +
37 85 default Long selectCountBySessionId(Long sessionId) {
38 86 return selectCount(Ai1SessionMessageDO::getSessionId, sessionId);
39 87 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/Ai1RedisKeyConstants.java
... ... @@ -39,4 +39,31 @@ public interface Ai1RedisKeyConstants {
39 39 */
40 40 String SESSION_RESULT_STREAM = "ai1:session:result:%d";
41 41  
  42 + /**
  43 + * 会话停止意图,按助手消息编号隔离;覆盖排队中、其他节点执行中的任务
  44 + * <p>
  45 + * KEY 格式:ai1:session:stop:{messageId}
  46 + * VALUE 数据类型:String,过期时间与结果流一致
  47 + */
  48 + String SESSION_STOP = "ai1:session:stop:%d";
  49 +
  50 + /**
  51 + * 会话停止通知,消息内容为助手消息编号,用于即时唤醒执行节点
  52 + */
  53 + String SESSION_STOP_TOPIC = "ai1:session:stop";
  54 +
  55 + /**
  56 + * 长期记忆归档锁,按租户互斥;使用 Redisson 看门狗续期
  57 + * <p>
  58 + * KEY 格式:ai1:memory:archive-lock:{tenantId}
  59 + */
  60 + String MEMORY_ARCHIVE_LOCK = "ai1:memory:archive-lock:%s";
  61 +
  62 + /**
  63 + * 长期记忆快照锁:消息落库事务与归档读取互斥,避免低编号消息晚提交后落在水位线之前
  64 + * <p>
  65 + * KEY 格式:ai1:memory:snapshot-lock:{agentId}:{userId};Agent 编号全局唯一,KEY 不带 tenantId
  66 + */
  67 + String MEMORY_SNAPSHOT_LOCK = "ai1:memory:snapshot-lock:%d:%d";
  68 +
42 69 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/dal/redis/memory/Ai1MemoryLockRedisDAO.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.dal.redis.memory;
  2 +
  3 +import jakarta.annotation.Resource;
  4 +import org.redisson.api.RLock;
  5 +import org.redisson.api.RedissonClient;
  6 +import org.springframework.stereotype.Repository;
  7 +
  8 +import java.util.function.Supplier;
  9 +
  10 +import static cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants.MEMORY_ARCHIVE_LOCK;
  11 +import static cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants.MEMORY_SNAPSHOT_LOCK;
  12 +
  13 +/**
  14 + * AI1 长期记忆的锁 Redis DAO
  15 + *
  16 + * 不设置租期,由 Redisson 看门狗续期
  17 + *
  18 + * @author 芋道源码
  19 + */
  20 +@Repository
  21 +public class Ai1MemoryLockRedisDAO {
  22 +
  23 + @Resource
  24 + private RedissonClient redissonClient;
  25 +
  26 + /**
  27 + * 按租户尝试获取归档锁并执行;其他节点归档中时不等待、不执行
  28 + *
  29 + * @param tenantId 租户编号
  30 + * @param supplier 归档逻辑
  31 + * @return 执行结果;未获取到锁时返回 null
  32 + */
  33 + public <V> V tryLockArchive(Long tenantId, Supplier<V> supplier) {
  34 + RLock lock = redissonClient.getLock(formatArchiveKey(tenantId));
  35 + if (!lock.tryLock()) {
  36 + return null;
  37 + }
  38 + try {
  39 + return supplier.get();
  40 + } finally {
  41 + if (lock.isHeldByCurrentThread()) {
  42 + lock.unlock();
  43 + }
  44 + }
  45 + }
  46 +
  47 + /**
  48 + * 获取用户在 Agent 下的快照锁并执行:消息落库事务与归档读取新消息互斥
  49 + *
  50 + * @param agentId Agent 编号
  51 + * @param userId 用户编号
  52 + * @param supplier 加锁执行的逻辑
  53 + * @return 执行结果
  54 + */
  55 + public <V> V lockSnapshot(Long agentId, Long userId, Supplier<V> supplier) {
  56 + RLock lock = redissonClient.getLock(formatSnapshotKey(agentId, userId));
  57 + lock.lock();
  58 + try {
  59 + return supplier.get();
  60 + } finally {
  61 + lock.unlock();
  62 + }
  63 + }
  64 +
  65 + private static String formatArchiveKey(Long tenantId) {
  66 + return String.format(MEMORY_ARCHIVE_LOCK, tenantId);
  67 + }
  68 +
  69 + private static String formatSnapshotKey(Long agentId, Long userId) {
  70 + return String.format(MEMORY_SNAPSHOT_LOCK, agentId, userId);
  71 + }
  72 +
  73 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/enums/Ai1ErrorCodeConstants.java
... ... @@ -77,4 +77,7 @@ public interface Ai1ErrorCodeConstants {
77 77 // ========== 会话消息 1-041-009-000 ==========
78 78 ErrorCode SESSION_MESSAGE_NOT_EXISTS = new ErrorCode(1_041_009_000, "会话消息不存在");
79 79  
  80 + // ========== 长期记忆 1-041-010-000 ==========
  81 + ErrorCode MEMORY_NOT_EXISTS = new ErrorCode(1_041_010_000, "长期记忆不存在");
  82 +
80 83 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/framework/ai/config/YudaoAi1Properties.java
... ... @@ -27,6 +27,11 @@ public class YudaoAi1Properties {
27 27 */
28 28 private Skill skill = new Skill();
29 29  
  30 + /**
  31 + * 长期记忆
  32 + */
  33 + private Memory memory = new Memory();
  34 +
30 35 @Data
31 36 public static class Milvus {
32 37  
... ... @@ -143,4 +148,19 @@ public class YudaoAi1Properties {
143 148  
144 149 }
145 150  
  151 + @Data
  152 + public static class Memory {
  153 +
  154 + /**
  155 + * 默认记忆正文最大长度(字符)
  156 + */
  157 + public static final int DEFAULT_MAX_LENGTH = 2000;
  158 +
  159 + /**
  160 + * 记忆正文最大长度(字符),超出时截断
  161 + */
  162 + private Integer maxLength = DEFAULT_MAX_LENGTH;
  163 +
  164 + }
  165 +
146 166 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/llm/Ai1LlmChatTool.java
... ... @@ -90,15 +90,15 @@ public class Ai1LlmChatTool {
90 90 // 2.1 思考过程:Spring AI 1.x 的 reasoning_content 是增量分片,直接转发(重复文本也可能是有效增量)
91 91 Object reasoningValue = generation.getOutput().getMetadata().get(METADATA_REASONING);
92 92 String delta = reasoningValue instanceof String value ? value : null;
93   - if (StrUtil.isNotBlank(delta)) {
  93 + if (StrUtil.isNotEmpty(delta)) {
94 94 thinkingText.append(delta);
95 95 if (onThinking != null) {
96 96 onThinking.accept(delta);
97 97 }
98 98 }
99   - // 2.2 回复内容
  99 + // 2.2 回复内容:只跳过空串,空格、换行分片也是有效内容
100 100 String chunk = generation.getOutput().getText();
101   - if (StrUtil.isNotBlank(chunk)) {
  101 + if (StrUtil.isNotEmpty(chunk)) {
102 102 contentText.append(chunk);
103 103 if (onContent != null) {
104 104 onContent.accept(chunk);
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveTool.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.memory;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.util.ObjUtil;
  5 +import cn.hutool.core.util.StrUtil;
  6 +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
  7 +import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
  8 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  9 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  10 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
  11 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
  12 +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO;
  13 +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
  14 +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
  15 +import cn.iocoder.yudao.module.ai1.harness.session.Ai1SessionStreamTool;
  16 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
  17 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
  18 +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService;
  19 +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
  20 +import jakarta.annotation.Resource;
  21 +import lombok.extern.slf4j.Slf4j;
  22 +import org.springframework.stereotype.Component;
  23 +
  24 +import java.util.Collections;
  25 +import java.util.List;
  26 +
  27 +import static cn.iocoder.yudao.framework.common.util.collection.CollectionUtils.convertList;
  28 +
  29 +/**
  30 + * AI1 长期记忆归档工具(harness 运行时)
  31 + *
  32 + * 【归档流程】由 Ai1MemoryArchiveJob 按租户调用 {@link #archive}
  33 + * 1. 竞争:按租户 tryLock 归档锁,抢不到直接跳过
  34 + * 2. 检测:遍历开启且开启记忆的 Agent,对其下每个用户,读取水位线 lastMessageId 之后的已完成消息
  35 + * 3. 生成:「历史记忆 + 新消息」经模型生成最新记忆,写入并推进水位线
  36 + *
  37 + * 【关键约定】
  38 + * 1. 新消息只取「安全前缀」:截止到水位线之后第一条生成中的消息,避免水位线越过尚未完成的回复
  39 + * 2. 超时且无 worker 执行的孤儿生成先按中断收尾,避免永久阻塞安全前缀
  40 + * 3. 单次最多归档 {@link #BATCH_SIZE} 条,剩余的由下次执行继续
  41 + * 4. 写入前校验:记忆被人工修改、删除,或 Agent 已删除时放弃本次结果
  42 + * 5. 锁顺序固定为「归档锁 → 消息执行锁(tryLock)」,快照锁只用于读取新消息,持有期间不收尾、不调用模型
  43 + *
  44 + * @author 芋道源码
  45 + */
  46 +@Component
  47 +@Slf4j
  48 +public class Ai1MemoryArchiveTool {
  49 +
  50 + /**
  51 + * 单个用户单次归档的最大消息数
  52 + */
  53 + static final int BATCH_SIZE = 100;
  54 +
  55 + @Resource
  56 + private Ai1MemoryLockRedisDAO memoryLockRedisDAO;
  57 + @Resource
  58 + private YudaoAi1Properties ai1Properties;
  59 +
  60 + @Resource
  61 + private Ai1AgentService agentService;
  62 + @Resource
  63 + private Ai1SessionService sessionService;
  64 + @Resource
  65 + private Ai1SessionMessageService sessionMessageService;
  66 + @Resource
  67 + private Ai1MemoryService memoryService;
  68 +
  69 + @Resource
  70 + private Ai1MemoryTool memoryTool;
  71 + @Resource
  72 + private Ai1SessionStreamTool sessionStreamTool;
  73 +
  74 + /**
  75 + * 归档当前租户下全部开启记忆的 Agent
  76 + *
  77 + * @return 归档成功的用户数
  78 + */
  79 + public int archive() {
  80 + // 1. 按租户竞争归档锁,其他节点归档中时直接跳过
  81 + Integer count = memoryLockRedisDAO.tryLockArchive(TenantContextHolder.getTenantId(), this::archiveAll);
  82 + if (count == null) {
  83 + log.info("[archive][租户({}) 归档进行中,跳过本次]", TenantContextHolder.getTenantId());
  84 + return 0;
  85 + }
  86 + return count;
  87 + }
  88 +
  89 + private int archiveAll() {
  90 + // 2. 逐个 Agent、逐个用户归档;单个用户失败不影响其他用户
  91 + int count = 0;
  92 + int maxLength = ai1Properties.getMemory().getMaxLength();
  93 + for (Ai1AgentDO agent : agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())) {
  94 + if (!Boolean.TRUE.equals(agent.getMemoryEnabled())) {
  95 + continue;
  96 + }
  97 + for (Long userId : sessionService.getUserIdListByAgentId(agent.getId())) {
  98 + if (Thread.currentThread().isInterrupted()) {
  99 + return count;
  100 + }
  101 + try {
  102 + if (archiveOne(agent, userId, maxLength)) {
  103 + count++;
  104 + }
  105 + } catch (Exception e) {
  106 + log.warn("[archive][Agent({}) 用户({}) 归档失败]", agent.getId(), userId, e);
  107 + }
  108 + }
  109 + }
  110 + return count;
  111 + }
  112 +
  113 + /**
  114 + * 归档单个用户:读取新消息 → 生成记忆 → 校验后写入并推进水位线
  115 + *
  116 + * @return 是否写入成功
  117 + */
  118 + boolean archiveOne(Ai1AgentDO agent, Long userId, int maxLength) {
  119 + // 1.1 读取水位线,并收尾水位线之后超时的孤儿生成
  120 + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(agent.getId(), userId);
  121 + Long lastMessageId = memory != null ? ObjUtil.defaultIfNull(memory.getLastMessageId(), 0L) : 0L;
  122 + failStaleMessages(agent.getId(), userId, lastMessageId);
  123 + // 1.2 读取水位线之后的新消息,无则跳过
  124 + List<Ai1SessionMessageDO> messages = getNewMessageList(agent.getId(), userId, lastMessageId);
  125 + if (CollUtil.isEmpty(messages)) {
  126 + return false;
  127 + }
  128 +
  129 + // 2. 生成最新记忆
  130 + String content = memoryTool.generate(agent, memory != null ? memory.getContent() : null, messages, maxLength);
  131 + if (StrUtil.isBlank(content)) {
  132 + return false;
  133 + }
  134 +
  135 + // 3.1 首次归档(含删除后重新生成):Agent 已删除,或记忆已存在时放弃,由下次执行基于新记忆继续
  136 + Long newLastMessageId = CollUtil.getLast(messages).getId();
  137 + if (memory == null) {
  138 + if (!memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(agent.getId()).userId(userId)
  139 + .content(content).lastMessageId(newLastMessageId).build())) {
  140 + log.info("[archiveOne][Agent({}) 用户({}) Agent 已删除或记忆已存在,放弃本次结果]", agent.getId(), userId);
  141 + return false;
  142 + }
  143 + return true;
  144 + }
  145 + // 3.2 再次归档:记忆被人工修改或删除时放弃
  146 + if (!memoryService.updateMemoryIfUnchanged(memory, content, newLastMessageId)) {
  147 + log.info("[archiveOne][Agent({}) 用户({}) 记忆已变更,放弃本次结果]", agent.getId(), userId);
  148 + return false;
  149 + }
  150 + return true;
  151 + }
  152 +
  153 + /**
  154 + * 收尾水位线之后超时的孤儿生成,直到遇到仍在执行或未超时的生成
  155 + *
  156 + * 在快照锁外执行:收尾需抢消息执行锁并读写结果流,避免持有快照锁期间阻塞发送消息
  157 + */
  158 + private void failStaleMessages(Long agentId, Long userId, Long lastMessageId) {
  159 + List<Ai1SessionDO> sessions = sessionService.getSessionListByAgentIdAndUserId(agentId, userId);
  160 + if (CollUtil.isEmpty(sessions)) {
  161 + return;
  162 + }
  163 + List<Long> sessionIds = convertList(sessions, Ai1SessionDO::getId);
  164 + for (int i = 0; i < BATCH_SIZE; i++) {
  165 + Ai1SessionMessageDO generatingMessage = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(
  166 + sessionIds, Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId);
  167 + if (generatingMessage == null || !sessionStreamTool.failStaleMessage(generatingMessage)) {
  168 + return;
  169 + }
  170 + }
  171 + }
  172 +
  173 + /**
  174 + * 获得水位线之后的新消息:截止到第一条生成中的消息,按编号升序,最多 {@link #BATCH_SIZE} 条
  175 + *
  176 + * 持有与消息创建相同的快照锁读取:创建在锁内提交,读取时不存在编号更小、尚未提交的消息
  177 + */
  178 + private List<Ai1SessionMessageDO> getNewMessageList(Long agentId, Long userId, Long lastMessageId) {
  179 + return memoryLockRedisDAO.lockSnapshot(agentId, userId, () -> {
  180 + List<Ai1SessionDO> sessions = sessionService.getSessionListByAgentIdAndUserId(agentId, userId);
  181 + if (CollUtil.isEmpty(sessions)) {
  182 + return Collections.emptyList();
  183 + }
  184 + List<Long> sessionIds = convertList(sessions, Ai1SessionDO::getId);
  185 + Ai1SessionMessageDO generatingMessage = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(
  186 + sessionIds, Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId);
  187 + return sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds,
  188 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), lastMessageId,
  189 + generatingMessage != null ? generatingMessage.getId() : null, BATCH_SIZE);
  190 + });
  191 + }
  192 +
  193 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryTool.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.memory;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.util.StrUtil;
  5 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  6 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  7 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
  8 +import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
  9 +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
  10 +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
  11 +import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool;
  12 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
  13 +import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService;
  14 +import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO;
  15 +import cn.iocoder.yudao.module.ai1.util.Ai1Utils;
  16 +import jakarta.annotation.Resource;
  17 +import org.springframework.stereotype.Component;
  18 +
  19 +import java.util.Collections;
  20 +import java.util.List;
  21 +
  22 +import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
  23 +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MODEL_TYPE_NOT_CHAT;
  24 +
  25 +/**
  26 + * AI1 长期记忆工具(harness 运行时)
  27 + *
  28 + * 1. {@link #loadMemory}:读取用户在 Agent 下的记忆
  29 + * 2. {@link #mergePrompt}:把记忆追加到系统指令,供对话注入
  30 + * 3. {@link #generate}:以「历史记忆 + 最新对话」为输入,经模型按固定维度生成最新记忆
  31 + *
  32 + * 记忆的持久化由归档 {@link Ai1MemoryArchiveTool} 完成
  33 + *
  34 + * @author 芋道源码
  35 + */
  36 +@Component
  37 +public class Ai1MemoryTool {
  38 +
  39 + /**
  40 + * 记忆维度小节标题
  41 + */
  42 + private static final List<String> MEMORY_SECTIONS = List.of("用户画像", "偏好习惯", "关键事实", "目标与意图", "近期进展");
  43 +
  44 + /**
  45 + * 记忆维度与生成约束(系统指令,{maxLen} 为长度上限占位符)
  46 + */
  47 + private static final String MEMORY_SYSTEM_PROMPT = """
  48 + 你是长期记忆归档助手。请基于【历史记忆】与【最新对话】,按下列维度抽取并更新用户长期记忆。
  49 +
  50 + 输出格式(务必严格遵守):
  51 + 1. 只输出 Markdown 记忆正文,不要任何解释、前言或代码块围栏;
  52 + 2. 每个维度标题独占一行,形如“## 维度名”(井号后有且仅有一个空格);
  53 + 3. 每个要点独占一行,以“- ”开头;
  54 + 4. 标题与要点之间、要点与要点之间都必须换行,禁止把标题与内容写在同一行;
  55 + 5. 不同维度之间空一行。
  56 +
  57 + 维度(无内容的维度整节省略):
  58 + ## 用户画像
  59 + - 用户的身份、称呼、职业、所在地等基本画像
  60 + ## 偏好习惯
  61 + - 用户的沟通偏好、回答风格、常用语言与工具等
  62 + ## 关键事实
  63 + - 用户明确告知、需要长期记住的事实性信息
  64 + ## 目标与意图
  65 + - 用户的核心诉求、长期目标与待办事项
  66 + ## 近期进展
  67 + - 最近对话中已推进的事项、结论与下一步
  68 +
  69 + 要求:合并历史记忆与新对话信息,删除重复与已过时内容,保持精炼;总长度不超过 {maxLen} 字。""";
  70 +
  71 + private static final String PLACEHOLDER_MAX_LEN = "{maxLen}";
  72 +
  73 + /**
  74 + * 合并进系统指令时,记忆正文前的引导语
  75 + */
  76 + private static final String MEMORY_PROMPT_PREFIX = "\n\n以下是该用户的历史长期记忆(Markdown),请在回答时结合这些信息;"
  77 + + "若与用户当前表述冲突,以当前表述为准:\n";
  78 +
  79 + @Resource
  80 + private Ai1MemoryService memoryService;
  81 + @Resource
  82 + private Ai1ModelService modelService;
  83 +
  84 + @Resource
  85 + private Ai1LlmChatTool llmChatTool;
  86 +
  87 + /**
  88 + * 读取用户在 Agent 下的长期记忆
  89 + *
  90 + * @param agentId Agent 编号
  91 + * @param userId 用户编号
  92 + * @return 记忆正文;不存在时返回 null
  93 + */
  94 + public String loadMemory(Long agentId, Long userId) {
  95 + if (agentId == null || userId == null) {
  96 + return null;
  97 + }
  98 + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(agentId, userId);
  99 + return memory != null ? memory.getContent() : null;
  100 + }
  101 +
  102 + /**
  103 + * 把长期记忆合并进系统指令;记忆为空时原样返回
  104 + *
  105 + * @param basePrompt 系统指令
  106 + * @param memory 记忆正文
  107 + * @return 合并后的系统指令
  108 + */
  109 + public String mergePrompt(String basePrompt, String memory) {
  110 + String prompt = StrUtil.nullToEmpty(basePrompt);
  111 + if (StrUtil.isBlank(memory)) {
  112 + return prompt;
  113 + }
  114 + return prompt + MEMORY_PROMPT_PREFIX + memory.trim();
  115 + }
  116 +
  117 + /**
  118 + * 经模型生成最新记忆:不装配历史、工具、RAG,也不注入已有记忆
  119 + *
  120 + * @param agent 记忆所属 Agent,使用其对话模型
  121 + * @param oldMemory 历史记忆,可为空
  122 + * @param newMessages 自上次归档以来的新消息,按编号升序
  123 + * @param maxLength 记忆长度上限(字符)
  124 + * @return 最新记忆;模型无输出时返回空
  125 + */
  126 + public String generate(Ai1AgentDO agent, String oldMemory, List<Ai1SessionMessageDO> newMessages, int maxLength) {
  127 + // 1. 校验对话模型
  128 + Ai1ModelRespBO model = modelService.getModelRespBO(agent.getProviderId(), agent.getModelId());
  129 + if (!Ai1ModelTypeEnum.isChat(model.getModelType())) {
  130 + throw exception(MODEL_TYPE_NOT_CHAT);
  131 + }
  132 +
  133 + // 2. 单轮会话,聚合为完整文本
  134 + int limit = maxLength > 0 ? maxLength : YudaoAi1Properties.Memory.DEFAULT_MAX_LENGTH;
  135 + String systemPrompt = MEMORY_SYSTEM_PROMPT.replace(PLACEHOLDER_MAX_LEN, String.valueOf(limit));
  136 + Ai1LlmChatTool.ChatText chatText = llmChatTool.chat(model, systemPrompt, Collections.emptyList(),
  137 + buildUserPrompt(oldMemory, newMessages), Collections.emptyList(), Collections.emptyList(),
  138 + null, null, null);
  139 +
  140 + // 3. 规范化格式并截断
  141 + return Ai1Utils.limitLength(Ai1Utils.normalizeMemoryMarkdown(chatText != null ? chatText.getContent() : null,
  142 + MEMORY_SECTIONS), limit);
  143 + }
  144 +
  145 + /**
  146 + * 组装归档的用户提示:历史记忆 + 最新对话
  147 + */
  148 + private static String buildUserPrompt(String oldMemory, List<Ai1SessionMessageDO> newMessages) {
  149 + StringBuilder prompt = new StringBuilder();
  150 + prompt.append("【历史记忆】\n");
  151 + prompt.append(StrUtil.isBlank(oldMemory) ? "(暂无)" : oldMemory.trim()).append("\n\n");
  152 + prompt.append("【最新对话】\n");
  153 + if (CollUtil.isNotEmpty(newMessages)) {
  154 + for (Ai1SessionMessageDO message : newMessages) {
  155 + String role = Ai1SessionMessageRoleEnum.isAssistant(message.getRole()) ? "助手" : "用户";
  156 + prompt.append(role).append(":").append(StrUtil.nullToEmpty(message.getContent())).append("\n");
  157 + }
  158 + }
  159 + return prompt.toString();
  160 + }
  161 +
  162 +}
... ...
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;
6 6 import cn.hutool.core.thread.ExecutorBuilder;
7 7 import cn.hutool.core.thread.ThreadUtil;
8 8 import cn.hutool.core.util.NumberUtil;
  9 +import cn.hutool.core.util.ObjUtil;
9 10 import cn.hutool.core.util.StrUtil;
10 11 import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
11 12 import cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil;
... ... @@ -13,6 +14,7 @@ import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
13 14 import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
14 15 import cn.iocoder.yudao.framework.tenant.core.util.TenantUtils;
15 16 import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  17 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
16 18 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
17 19 import cn.iocoder.yudao.module.ai1.dal.dataobject.knowledge.Ai1KnowledgeBaseDO;
18 20 import cn.iocoder.yudao.module.ai1.dal.redis.Ai1RedisKeyConstants;
... ... @@ -22,6 +24,8 @@ import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
22 24 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
23 25 import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
24 26 import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
  27 +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool;
  28 +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker;
25 29 import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory;
26 30 import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
27 31 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
... ... @@ -33,12 +37,15 @@ import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool;
33 37 import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory;
34 38 import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool;
35 39 import jakarta.annotation.Resource;
  40 +import lombok.Getter;
36 41 import lombok.extern.slf4j.Slf4j;
37 42 import org.springframework.ai.chat.client.advisor.api.Advisor;
38 43 import org.springframework.context.SmartLifecycle;
39 44 import org.springframework.context.annotation.Lazy;
  45 +import org.redisson.api.RBucket;
40 46 import org.redisson.api.RStream;
41 47 import org.redisson.api.RLock;
  48 +import org.redisson.api.RTopic;
42 49 import org.redisson.api.RedissonClient;
43 50 import org.redisson.api.stream.*;
44 51 import org.redisson.client.codec.StringCodec;
... ... @@ -46,10 +53,14 @@ import org.springframework.stereotype.Component;
46 53 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
47 54  
48 55 import java.time.Duration;
  56 +import java.time.LocalDateTime;
49 57 import java.util.*;
  58 +import java.util.concurrent.CancellationException;
50 59 import java.util.concurrent.ConcurrentHashMap;
51 60 import java.util.concurrent.ExecutorService;
  61 +import java.util.concurrent.Executors;
52 62 import java.util.concurrent.RejectedExecutionException;
  63 +import java.util.concurrent.ScheduledExecutorService;
53 64 import java.util.concurrent.TimeUnit;
54 65 import java.util.concurrent.atomic.AtomicBoolean;
55 66 import java.util.function.Consumer;
... ... @@ -62,6 +73,7 @@ import java.util.function.Consumer;
62 73 * 2. 生成(worker 侧,{@link #start} 拉起的消费线程):{@link #consumeLoop} → {@link #handleTask} → {@link #generate};
63 74 * 增量写入结果流,结束后回填助手消息并 XACK
64 75 * 3. 下发与续传(连接侧):{@link #streamResult} 从结果流 XREAD 转发为 SSE;携带 lastEventId 时从断点之后重放,不重新生成
  76 + * 4. 停止(请求侧):{@link #cancel} 按助手消息写入停止意图并广播,执行节点中断生成;排队中的任务直接收尾
65 77 *
66 78 * 【关键约定】
67 79 * 1. 助手消息编号同时作为结果流标识;生成与连接解耦,任意节点都可以转发、续传,worker 可独立扩容
... ... @@ -86,6 +98,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
86 98 * 生成 worker 线程名前缀,完整线程名为「前缀 + 消费者名」
87 99 */
88 100 private static final String WORKER_THREAD_NAME_PREFIX = "ai1-session-";
  101 + /**
  102 + * 停止意图检查线程名
  103 + */
  104 + private static final String STOP_THREAD_NAME = "ai1-session-stop-";
  105 + /**
  106 + * Shell 进程回收线程名前缀,完整线程名为「前缀 + worker 线程名」
  107 + */
  108 + private static final String REAPER_THREAD_NAME_PREFIX = "ai1-session-reap-";
89 109  
90 110 /**
91 111 * 任务队列最大长度(近似裁剪)
... ... @@ -120,6 +140,10 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
120 140 */
121 141 private static final long WORKER_JOIN_TIMEOUT_MILLIS = 3000;
122 142 /**
  143 + * worker 停止时,等待全部 Shell 进程回收线程退出的最长时长,单位:毫秒;单个进程回收最多等待 5s
  144 + */
  145 + private static final long REAPER_JOIN_TIMEOUT_MILLIS = 6000;
  146 + /**
123 147 * 结果流连续读取失败的最大次数,达到后中断转发
124 148 */
125 149 private static final int READ_ERROR_MAX_COUNT = 3;
... ... @@ -129,6 +153,11 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
129 153 private static final long READ_ERROR_BACKOFF_MILLIS = 500;
130 154  
131 155 /**
  156 + * 停止意图的检查间隔,单位:毫秒;兜底停止通知在订阅断开期间丢失
  157 + */
  158 + private static final long STOP_CHECK_INTERVAL_MILLIS = 1000;
  159 +
  160 + /**
132 161 * 结果流丢失判定:转发到连接超时前该时长,结果流仍不存在时,判定任务丢失(预留两个阻塞读窗口,保证 error 能在连接超时前送达)
133 162 */
134 163 private static final long LOST_CHECK_ADVANCE_MILLIS = READ_BLOCK_MILLIS * 2;
... ... @@ -143,6 +172,21 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
143 172 */
144 173 static final String INTERRUPTED_MESSAGE = "生成任务已中断,请重新发送";
145 174  
  175 + /**
  176 + * 人工终止、生成失败时追加到正文的提示;失败不拼接原始异常
  177 + */
  178 + static final String TERMINATED_NOTICE = "本条回复已被人工终止";
  179 + static final String FAILED_NOTICE = "本条回复生成失败";
  180 + /**
  181 + * 正文与提示之间的分隔
  182 + */
  183 + private static final String NOTICE_PREFIX = "\n\n> ";
  184 +
  185 + /**
  186 + * SSE 转发异常的提示;只是连接异常,后台生成不受影响,可刷新后续传
  187 + */
  188 + static final String TRANSPORT_ERROR_MESSAGE = "会话连接异常,请刷新后重试";
  189 +
146 190 private static final String FIELD_MESSAGE_ID = "messageId";
147 191 private static final String FIELD_AGENT_ID = "agentId";
148 192 private static final String FIELD_SESSION_ID = "sessionId";
... ... @@ -173,6 +217,8 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
173 217 private Ai1LlmChatTool llmChatTool;
174 218 @Resource
175 219 private Ai1RagTool ragTool;
  220 + @Resource
  221 + private Ai1MemoryTool memoryTool;
176 222  
177 223 @Resource
178 224 private Ai1McpToolFactory mcpToolFactory;
... ... @@ -184,6 +230,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
184 230 */
185 231 private final Map<Long, Long> resultExpireTimes = new ConcurrentHashMap<>();
186 232 /**
  233 + * 本节点执行中的生成:助手消息编号 → 运行上下文,供停止时中断生成线程、回收 Shell 进程
  234 + */
  235 + private final Map<Long, RunContext> runContexts = new ConcurrentHashMap<>();
  236 + /**
  237 + * 未结束的 Shell 进程回收线程:生成结束后可能仍在回收,停止时统一等待;线程结束时自行移除
  238 + */
  239 + private final Set<Thread> reapers = ConcurrentHashMap.newKeySet();
  240 + /**
187 241 * 实例标识,消费者名按「实例 + 序号」隔离
188 242 */
189 243 private final String instanceId = UUID.randomUUID().toString().substring(0, 8);
... ... @@ -199,6 +253,14 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
199 253 * SSE 转发线程池:单个转发阻塞于 XREAD 直到流结束,用有界池控制并发
200 254 */
201 255 private ExecutorService sseExecutor;
  256 + /**
  257 + * 停止通知的订阅编号
  258 + */
  259 + private Integer stopListenerId;
  260 + /**
  261 + * 停止意图的定时检查线程
  262 + */
  263 + private ScheduledExecutorService stopChecker;
202 264  
203 265 // ==================== 生命周期 ====================
204 266  
... ... @@ -224,7 +286,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
224 286 // 2. 初始化任务消费组
225 287 initTaskGroup();
226 288  
227   - // 3. 拉起 worker 消费线程,消费者名按「实例 + 序号」隔离
  289 + // 3. 订阅停止通知,即时中断本节点的生成;通知不持久,另定时检查停止意图兜底
  290 + // 中断交给检查线程执行,避免 Redisson 订阅线程等待增量回调
  291 + stopChecker = Executors.newSingleThreadScheduledExecutor(ThreadUtil.newNamedThreadFactory(STOP_THREAD_NAME, true));
  292 + stopChecker.scheduleWithFixedDelay(this::checkStopRequests,
  293 + STOP_CHECK_INTERVAL_MILLIS, STOP_CHECK_INTERVAL_MILLIS, TimeUnit.MILLISECONDS);
  294 + stopListenerId = getStopTopic().addListener(String.class,
  295 + (channel, messageId) -> stopChecker.execute(() -> cancelRunContext(Convert.toLong(messageId))));
  296 +
  297 + // 4. 拉起 worker 消费线程,消费者名按「实例 + 序号」隔离
228 298 running = true;
229 299 int workerCount = ai1Properties.getSession().getWorker().getCount();
230 300 for (int i = 0; i < workerCount; i++) {
... ... @@ -237,11 +307,33 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
237 307 }
238 308  
239 309 /**
240   - * 停止:置停止标志,等待各消费线程退出(每个最多 3s),最后关闭 SSE 转发线程池
  310 + * 停止:置停止标志,取消停止订阅,中断本节点执行中的生成,等待各消费线程退出(每个最多 3s)、
  311 + * Shell 进程回收(共最多 6s),最后关闭 SSE 转发线程池
  312 + *
  313 + * 中断的生成保留部分内容并标记失败,不会被其他节点认领后重新生成
241 314 */
242 315 @Override
243 316 public void stop() {
  317 + // 1. 先置停止标志,再在 runContexts 锁内取快照:与 registerRunContext 互斥,
  318 + // 之后登记的上下文由登记方自行中断,不会漏掉阻塞读取刚返回、尚未登记的任务
244 319 running = false;
  320 + if (stopListenerId != null) {
  321 + try {
  322 + getStopTopic().removeListener(stopListenerId);
  323 + } catch (Exception e) {
  324 + log.warn("[stop][停止通知取消订阅失败]", e);
  325 + }
  326 + }
  327 + if (stopChecker != null) {
  328 + stopChecker.shutdownNow();
  329 + }
  330 + List<RunContext> contexts;
  331 + synchronized (runContexts) {
  332 + contexts = new ArrayList<>(runContexts.values());
  333 + }
  334 + // 中断只启动各自的回收线程,单个进程回收慢不影响其他上下文
  335 + contexts.forEach(context -> context.cancel(INTERRUPTED_MESSAGE));
  336 + // 2. 等待消费线程退出
245 337 for (Thread worker : workers) {
246 338 try {
247 339 worker.join(WORKER_JOIN_TIMEOUT_MILLIS);
... ... @@ -250,6 +342,19 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
250 342 }
251 343 }
252 344 workers.clear();
  345 + // 3. 等待全部未结束的 Shell 进程回收,包括此前人工停止、生成已结束的回收;各回收线程并行,共用一个截止时刻
  346 + long reaperDeadline = System.currentTimeMillis() + REAPER_JOIN_TIMEOUT_MILLIS;
  347 + try {
  348 + for (Thread reaper : reapers) {
  349 + long remaining = reaperDeadline - System.currentTimeMillis();
  350 + if (remaining <= 0) {
  351 + break;
  352 + }
  353 + reaper.join(remaining);
  354 + }
  355 + } catch (InterruptedException e) {
  356 + Thread.currentThread().interrupt();
  357 + }
253 358 if (sseExecutor != null) {
254 359 sseExecutor.shutdown();
255 360 }
... ... @@ -268,15 +373,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
268 373 try {
269 374 getTaskStream().createGroup(StreamCreateGroupArgs.name(Ai1RedisKeyConstants.SESSION_TASK_GROUP).id(StreamMessageId.ALL).makeStream());
270 375 } catch (Exception e) {
271   - if (!isBusyGroup(e)) {
  376 + if (!isRedisError(e, "BUSYGROUP")) {
272 377 log.warn("[initTaskGroup][Chat 任务消费组初始化失败]", e);
273 378 }
274 379 }
275 380 }
276 381  
277   - private static boolean isBusyGroup(Throwable e) {
  382 + private static boolean isRedisError(Throwable e, String code) {
278 383 for (Throwable cause = e; cause != null; cause = cause.getCause()) {
279   - if (StrUtil.contains(cause.getMessage(), "BUSYGROUP")) {
  384 + if (StrUtil.contains(cause.getMessage(), code)) {
280 385 return true;
281 386 }
282 387 }
... ... @@ -331,8 +436,9 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
331 436 emitter.complete();
332 437 }
333 438 } catch (Exception e) {
  439 + // 只是转发异常,不改动消息状态,也不按生成失败提示
334 440 log.warn("[open][助手消息({}) 会话流转发异常]", messageId, e);
335   - sendErrorAndComplete(emitter, e.getMessage());
  441 + sendErrorAndComplete(emitter, TRANSPORT_ERROR_MESSAGE);
336 442 }
337 443 });
338 444 } catch (RejectedExecutionException e) {
... ... @@ -354,6 +460,151 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
354 460 return emitter;
355 461 }
356 462  
  463 + // ==================== 停止 ====================
  464 +
  465 + /**
  466 + * 停止会话中生成中的助手消息
  467 + *
  468 + * 按助手消息写入停止意图并广播,覆盖其他节点执行中、排队中的任务
  469 + *
  470 + * @param sessionId 会话编号
  471 + */
  472 + public void cancel(Long sessionId) {
  473 + // 1. 按请求时刻定位生成中的助手消息;之后发送的新一轮不受影响
  474 + List<Long> messageIds = sessionMessageService.getGeneratingAssistantMessageIdList(sessionId);
  475 + Duration ttl = Duration.ofSeconds(ai1Properties.getSession().getStream().getTtlSeconds());
  476 + for (Long messageId : messageIds) {
  477 + // 2.1 写入停止意图:worker 执行前、执行中都会检查,与结果流的过期时间一致
  478 + getStopBucket(messageId).set("1", ttl);
  479 + // 2.2 广播停止通知,执行节点即时中断
  480 + getStopTopic().publish(String.valueOf(messageId));
  481 + // 2.3 没有 worker 执行时(排队中,或 worker 已退出),直接收尾
  482 + cancelQueuedTask(messageId);
  483 + }
  484 + }
  485 +
  486 + /**
  487 + * 收尾无人执行的任务:抢到执行锁说明没有 worker 在生成,保留已有结果并标记终止
  488 + *
  489 + * 任务记录仍在队列中,之后由 worker 按「已回填」补齐终态并确认,不再调用模型
  490 + */
  491 + private void cancelQueuedTask(Long messageId) {
  492 + RLock lock = redissonClient.getLock(String.format(Ai1RedisKeyConstants.SESSION_TASK_LOCK, messageId));
  493 + if (!lock.tryLock()) {
  494 + return;
  495 + }
  496 + try {
  497 + Ai1SessionMessageDO message = sessionMessageService.getSessionMessage(messageId);
  498 + if (message == null || !Ai1SessionMessageStatusEnum.isGenerating(message.getStatus())) {
  499 + return;
  500 + }
  501 + saveAssistantMessage(messageId, message.getSessionId(), withNotice(readPartialResult(messageId), TERMINATED_NOTICE),
  502 + Ai1SessionMessageStatusEnum.FAILED.getStatus());
  503 + appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), "");
  504 + } finally {
  505 + if (lock.isHeldByCurrentThread()) {
  506 + lock.unlock();
  507 + }
  508 + }
  509 + }
  510 +
  511 + /**
  512 + * 收尾超时的孤儿生成:任务记录被裁剪或丢失时,助手消息会一直停在生成中,阻塞长期记忆归档
  513 + *
  514 + * 超时阈值与超时任务认领一致,仍有任务记录的消息此前已被认领收尾;
  515 + * 先判定超时再抢执行锁,避免短暂持锁使 worker 把排队中的新任务误判为在途而跳过
  516 + *
  517 + * @param message 生成中的助手消息
  518 + * @return 是否已不在生成中;仍在执行或未超时时返回 false
  519 + * @throws IllegalStateException 收尾期间执行锁失效,由归档按单个用户失败处理
  520 + */
  521 + public boolean failStaleMessage(Ai1SessionMessageDO message) {
  522 + // 1. 未超时直接保留,不加锁
  523 + if (!isStaleMessage(message)) {
  524 + return false;
  525 + }
  526 + // 2. 抢不到执行锁说明仍有 worker 在生成
  527 + Long messageId = message.getId();
  528 + RLock lock = redissonClient.getLock(String.format(Ai1RedisKeyConstants.SESSION_TASK_LOCK, messageId));
  529 + if (!lock.tryLock()) {
  530 + return false;
  531 + }
  532 + long ownerThreadId = Thread.currentThread().threadId();
  533 + try {
  534 + // 3. 锁内重新读取,期间可能已被 worker 完成
  535 + Ai1SessionMessageDO current = sessionMessageService.getSessionMessage(messageId);
  536 + if (current == null || !Ai1SessionMessageStatusEnum.isGenerating(current.getStatus())) {
  537 + return true;
  538 + }
  539 + if (!isStaleMessage(current)) {
  540 + return false;
  541 + }
  542 + // 4. 与接管任务的收尾一致:保留部分结果并标记失败,再写入 error、done;
  543 + // 每次写入前校验执行权,锁失效后由新持有者处理,不再改动消息或结果流
  544 + log.warn("[failStaleMessage][助手消息({}) 生成超时且无 worker 执行,标记中断]", messageId);
  545 + Ai1LlmChatTool.ChatText partialText = readPartialResult(messageId);
  546 + checkTaskLock(lock, ownerThreadId);
  547 + saveAssistantMessage(messageId, current.getSessionId(), withNotice(partialText, INTERRUPTED_MESSAGE),
  548 + Ai1SessionMessageStatusEnum.FAILED.getStatus());
  549 + checkTaskLock(lock, ownerThreadId);
  550 + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE);
  551 + checkTaskLock(lock, ownerThreadId);
  552 + appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), "");
  553 + return true;
  554 + } finally {
  555 + if (lock.isHeldByThread(ownerThreadId)) {
  556 + lock.unlock();
  557 + }
  558 + }
  559 + }
  560 +
  561 + private boolean isStaleMessage(Ai1SessionMessageDO message) {
  562 + if (!Ai1SessionMessageRoleEnum.isAssistant(message.getRole()) || message.getCreateTime() == null) {
  563 + return false;
  564 + }
  565 + long staleMillis = ai1Properties.getSession().getStream().getTimeoutMs() + RECLAIM_IDLE_BUFFER_MILLIS;
  566 + return message.getCreateTime().isBefore(LocalDateTime.now().minus(Duration.ofMillis(staleMillis)));
  567 + }
  568 +
  569 + /**
  570 + * 登记运行上下文;节点已停机时立即中断
  571 + *
  572 + * 与 {@link #stop} 的快照在同一把锁内:要么被快照中断,要么在此看到停止标志
  573 + */
  574 + private void registerRunContext(Long messageId, RunContext context) {
  575 + synchronized (runContexts) {
  576 + runContexts.put(messageId, context);
  577 + if (!running) {
  578 + context.cancel(INTERRUPTED_MESSAGE);
  579 + }
  580 + }
  581 + }
  582 +
  583 + /**
  584 + * 中断本节点执行中的生成;不在本节点执行时忽略
  585 + */
  586 + private void cancelRunContext(Long messageId) {
  587 + RunContext context = messageId != null ? runContexts.get(messageId) : null;
  588 + if (context != null) {
  589 + context.cancel(TERMINATED_NOTICE);
  590 + }
  591 + }
  592 +
  593 + /**
  594 + * 定时检查本节点执行中的生成是否有停止意图,兜底订阅断开期间丢失的停止通知
  595 + */
  596 + private void checkStopRequests() {
  597 + runContexts.keySet().forEach(messageId -> {
  598 + try {
  599 + if (getStopBucket(messageId).isExists()) {
  600 + cancelRunContext(messageId);
  601 + }
  602 + } catch (Exception e) {
  603 + log.warn("[checkStopRequests][助手消息({}) 停止意图检查失败]", messageId, e);
  604 + }
  605 + });
  606 + }
  607 +
357 608 // ==================== worker 消费 ====================
358 609  
359 610 /**
... ... @@ -382,6 +633,10 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
382 633 return;
383 634 }
384 635 log.warn("[consumeLoop][Chat worker({}) 消费异常]", consumer, e);
  636 + // 任务队列被删除(如 Redis 清库)后消费组随之丢失,重建后继续消费
  637 + if (isRedisError(e, "NOGROUP")) {
  638 + initTaskGroup();
  639 + }
385 640 ThreadUtil.sleep(CONSUME_ERROR_BACKOFF_MILLIS);
386 641 }
387 642 }
... ... @@ -477,7 +732,8 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
477 732 Ai1LlmChatTool.ChatText partialText = readPartialResult(messageId);
478 733 // 1.3.2 回填部分结果并标记失败;每次写入前校验执行权,防止旧 worker 继续写入
479 734 checkTaskLock(lock, workerThreadId);
480   - saveAssistantMessage(messageId, sessionId, partialText, Ai1SessionMessageStatusEnum.FAILED.getStatus());
  735 + saveAssistantMessage(messageId, sessionId, withNotice(partialText, INTERRUPTED_MESSAGE),
  736 + Ai1SessionMessageStatusEnum.FAILED.getStatus());
481 737 // 1.3.3 写入错误提示和结束事件,续传连接可及时退出
482 738 checkTaskLock(lock, workerThreadId);
483 739 appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE);
... ... @@ -493,40 +749,81 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
493 749 Long agentId = Convert.toLong(fields.get(FIELD_AGENT_ID));
494 750 String content = fields.get(FIELD_CONTENT);
495 751  
496   - // 3. 生成并回填;增量回调可能在模型线程执行,因此按 worker 线程编号校验锁归属
  752 + // 3. 登记运行上下文,供停止时中断生成线程、回收 Shell 进程
  753 + RunContext context = new RunContext(Thread.currentThread(), skillToolFactory.newShellTracker(agentId), reapers);
  754 + registerRunContext(messageId, context);
  755 + // 增量回调可能在模型线程执行,因此按 worker 线程编号校验锁归属;停止后不再写入增量
497 756 StringBuilder contentText = new StringBuilder();
498 757 StringBuilder thinkingText = new StringBuilder();
499 758 Consumer<String> onThinking = delta -> {
500   - checkTaskLock(lock, workerThreadId);
501   - thinkingText.append(delta);
502   - appendResult(messageId, Ai1SessionStreamEventEnum.THINKING.getEvent(), delta);
  759 + synchronized (context) {
  760 + context.checkActive();
  761 + checkTaskLock(lock, workerThreadId);
  762 + thinkingText.append(delta);
  763 + appendResult(messageId, Ai1SessionStreamEventEnum.THINKING.getEvent(), delta);
  764 + }
503 765 };
504 766 Consumer<String> onContent = delta -> {
505   - checkTaskLock(lock, workerThreadId);
506   - contentText.append(delta);
507   - appendResult(messageId, Ai1SessionStreamEventEnum.MESSAGE.getEvent(), delta);
  767 + synchronized (context) {
  768 + context.checkActive();
  769 + checkTaskLock(lock, workerThreadId);
  770 + contentText.append(delta);
  771 + appendResult(messageId, Ai1SessionStreamEventEnum.MESSAGE.getEvent(), delta);
  772 + }
508 773 };
509   - // 3.1 标记生成开始,再调用模型并回填完整结果
  774 +
  775 + // 4. 标记生成开始,再调用模型
  776 + Ai1LlmChatTool.ChatText chatText = null;
  777 + Exception failure = null;
  778 + String stopNotice;
510 779 try {
  780 + // 4.1 节点已停机,或排队期间已被停止时,不再调用模型
  781 + context.checkActive();
  782 + if (getStopBucket(messageId).isExists()) {
  783 + context.cancel(TERMINATED_NOTICE);
  784 + }
  785 + context.checkActive();
  786 + // 4.2 标记生成已开始;长思考没有增量时,续传也不会误判结果流丢失
511 787 checkTaskLock(lock, workerThreadId);
512   - // 标记生成已开始;长思考没有增量时,续传也不会误判结果流丢失
513 788 appendResult(messageId, Ai1SessionStreamEventEnum.PING.getEvent(), "");
514   - Ai1LlmChatTool.ChatText chatText = generate(messageId, agentId, sessionId, content, onThinking, onContent);
515   - checkTaskLock(lock, workerThreadId);
516   - saveAssistantMessage(messageId, sessionId, chatText, Ai1SessionMessageStatusEnum.SUCCESS.getStatus());
  789 + chatText = generate(context, messageId, agentId, sessionId, content, onThinking, onContent);
517 790 } catch (Exception e) {
518   - // 失去锁的旧 worker 不得再写结果、终态或 XACK,留给新持有者处理
  791 + failure = e;
  792 + } finally {
  793 + // 4.3 结束上下文并清除中断标志,之后的停止不再生效
  794 + stopNotice = context.finish();
  795 + runContexts.remove(messageId, context);
  796 + context.getShellTracker().cancel();
  797 + }
  798 +
  799 + // 5. 先回填终态,再下发 error / done;每次写入前校验执行权
  800 + if (stopNotice != null) {
  801 + // 5.1 停止:保留已生成的内容并追加提示;人工终止只写 done,节点停机中断另写 error
  802 + log.info("[handleTask][助手消息({}) 生成已停止:{}]", messageId, stopNotice);
  803 + Ai1LlmChatTool.ChatText partialText = chatText != null ? chatText
  804 + : new Ai1LlmChatTool.ChatText(contentText.toString(), thinkingText.toString());
  805 + checkTaskLock(lock, workerThreadId);
  806 + saveAssistantMessage(messageId, sessionId, withNotice(partialText, stopNotice), Ai1SessionMessageStatusEnum.FAILED.getStatus());
  807 + if (!TERMINATED_NOTICE.equals(stopNotice)) {
  808 + checkTaskLock(lock, workerThreadId);
  809 + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), stopNotice);
  810 + }
  811 + } else if (failure != null) {
  812 + // 5.2 失败:失去锁的旧 worker 不得再写结果、终态或 XACK,留给新持有者处理
519 813 if (!lock.isHeldByThread(workerThreadId)) {
520 814 return;
521 815 }
522   - // 3.2 生成失败时下发错误,保留本轮已产生的增量并回填失败状态
523   - log.warn("[handleTask][助手消息({}) 生成失败]", messageId, e);
524   - appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), StrUtil.blankToDefault(e.getMessage(), "生成失败"));
  816 + log.warn("[handleTask][助手消息({}) 生成失败]", messageId, failure);
  817 + saveAssistantMessage(messageId, sessionId, withNotice(new Ai1LlmChatTool.ChatText(contentText.toString(),
  818 + thinkingText.toString()), FAILED_NOTICE), Ai1SessionMessageStatusEnum.FAILED.getStatus());
525 819 checkTaskLock(lock, workerThreadId);
526   - saveAssistantMessage(messageId, sessionId, new Ai1LlmChatTool.ChatText(contentText.toString(),
527   - thinkingText.toString()), Ai1SessionMessageStatusEnum.FAILED.getStatus());
  820 + appendResult(messageId, Ai1SessionStreamEventEnum.ERROR.getEvent(), FAILED_NOTICE);
  821 + } else {
  822 + // 5.3 成功:回填完整结果
  823 + checkTaskLock(lock, workerThreadId);
  824 + saveAssistantMessage(messageId, sessionId, chatText, Ai1SessionMessageStatusEnum.SUCCESS.getStatus());
528 825 }
529   - // 4. 回填后写终态并确认;失败时由后续接管清理,不能在 finally 中无条件确认
  826 + // 6. 回填后写终态并确认;失败时由后续接管清理,不能在 finally 中无条件确认
530 827 checkTaskLock(lock, workerThreadId);
531 828 appendResult(messageId, Ai1SessionStreamEventEnum.DONE.getEvent(), "");
532 829 checkTaskLock(lock, workerThreadId);
... ... @@ -546,11 +843,15 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
546 843 if (Ai1SessionStreamEventEnum.DONE.getEvent().equals(lastType)) {
547 844 return;
548 845 }
549   - // 2. 失败消息缺少 error 时补写错误提示,末尾已有 error 时不重复写入
  846 + // 2. 失败消息缺少 error 时,按正文中已回填的提示补写;人工终止只写 done,末尾已有 error 时不重复写入
550 847 if (Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())
551 848 && !Ai1SessionStreamEventEnum.ERROR.getEvent().equals(lastType)) {
552   - checkTaskLock(lock, workerThreadId);
553   - appendResult(message.getId(), Ai1SessionStreamEventEnum.ERROR.getEvent(), INTERRUPTED_MESSAGE);
  849 + String notice = getSavedNotice(message.getContent());
  850 + if (!TERMINATED_NOTICE.equals(notice)) {
  851 + checkTaskLock(lock, workerThreadId);
  852 + appendResult(message.getId(), Ai1SessionStreamEventEnum.ERROR.getEvent(),
  853 + ObjUtil.defaultIfNull(notice, INTERRUPTED_MESSAGE));
  854 + }
554 855 }
555 856 // 3. 补写 done,使续传连接能够结束;调用方在补齐后再确认任务
556 857 checkTaskLock(lock, workerThreadId);
... ... @@ -593,12 +894,12 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
593 894 // ==================== 生成编排 ====================
594 895  
595 896 /**
596   - * 执行一次生成:校验 Agent → 解析模型 → 装配历史、工具、RAG → 流式会话
  897 + * 执行一次生成:校验 Agent → 解析模型 → 装配历史、长期记忆、工具、RAG → 流式会话
597 898 *
598 899 * 历史只取当前助手占位之前最近 historyLimit 条,跳过未完成的助手占位,并移除末尾的当前提问(由 LLM 工具另行追加)
599 900 */
600 901 @SuppressWarnings("SequencedCollectionMethodCanBeUsed")
601   - private Ai1LlmChatTool.ChatText generate(Long messageId, Long agentId, Long sessionId, String content,
  902 + private Ai1LlmChatTool.ChatText generate(RunContext context, Long messageId, Long agentId, Long sessionId, String content,
602 903 Consumer<String> onThinking, Consumer<String> onContent) {
603 904 // 1. 生成可能发生在其他节点,重新校验 Agent 与模型;后台会话不要求 Agent 已发布
604 905 Ai1AgentDO agent = agentService.validateAgentExists(agentId);
... ... @@ -631,10 +932,13 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
631 932 String systemPrompt = StrUtil.isNotBlank(agent.getSystemPrompt()) ? agent.getSystemPrompt().trim()
632 933 : StrUtil.replace(ai1Properties.getSession().getDefaultSystemPrompt(),
633 934 YudaoAi1Properties.Session.AGENT_NAME_PLACEHOLDER, agent.getName());
  935 + // 开启长期记忆时,合入用户在该 Agent 下的记忆
  936 + systemPrompt = injectMemory(agent, sessionId, systemPrompt);
634 937  
635 938 // 4. 工具(MCP + SKILL)与 RAG Advisor;知识库不可用时降级跳过,并把提示写入思考流
  939 + // Shell 命令经运行上下文的跟踪器执行,停止或生成结束时统一回收
636 940 List<Object> tools = new ArrayList<>(mcpToolFactory.buildTools(agent));
637   - tools.addAll(skillToolFactory.buildTools(agent));
  941 + tools.addAll(skillToolFactory.buildTools(agent, context.getShellTracker()));
638 942 StringBuilder noticeText = new StringBuilder();
639 943 List<Advisor> advisors = new ArrayList<>();
640 944 for (Ai1KnowledgeBaseDO knowledgeBase : knowledgeBaseService.getKnowledgeBaseList(agent.getKnowledgeBaseIds())) {
... ... @@ -652,6 +956,7 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
652 956 }
653 957  
654 958 // 5. 流式会话;降级提示并入最终思考文本,保证刷新后内容一致
  959 + context.checkActive();
655 960 Ai1LlmChatTool.ChatText chatText = llmChatTool.chat(model, systemPrompt, histories, content, tools, advisors,
656 961 sessionId, onThinking, onContent);
657 962 if (!noticeText.isEmpty()) {
... ... @@ -661,6 +966,27 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
661 966 }
662 967  
663 968 /**
  969 + * 长期记忆注入
  970 + *
  971 + * 查询失败时降级为原系统指令,保证对话可用
  972 + */
  973 + private String injectMemory(Ai1AgentDO agent, Long sessionId, String systemPrompt) {
  974 + if (!Boolean.TRUE.equals(agent.getMemoryEnabled())) {
  975 + return systemPrompt;
  976 + }
  977 + try {
  978 + Ai1SessionDO session = sessionService.getSession(sessionId);
  979 + if (session == null) {
  980 + return systemPrompt;
  981 + }
  982 + return memoryTool.mergePrompt(systemPrompt, memoryTool.loadMemory(agent.getId(), session.getUserId()));
  983 + } catch (Exception e) {
  984 + log.warn("[injectMemory][Agent({}) 会话({}) 长期记忆注入失败,已降级跳过]", agent.getId(), sessionId, e);
  985 + return systemPrompt;
  986 + }
  987 + }
  988 +
  989 + /**
664 990 * 回填助手消息(占位记录)的内容、思考过程与生成状态,并刷新会话活跃时间
665 991 */
666 992 private void saveAssistantMessage(Long messageId, Long sessionId, Ai1LlmChatTool.ChatText chatText, Integer status) {
... ... @@ -830,6 +1156,33 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
830 1156 return redissonClient.getStream(String.format(Ai1RedisKeyConstants.SESSION_RESULT_STREAM, messageId), StringCodec.INSTANCE);
831 1157 }
832 1158  
  1159 + private RBucket<String> getStopBucket(Long messageId) {
  1160 + return redissonClient.getBucket(String.format(Ai1RedisKeyConstants.SESSION_STOP, messageId), StringCodec.INSTANCE);
  1161 + }
  1162 +
  1163 + private RTopic getStopTopic() {
  1164 + return redissonClient.getTopic(Ai1RedisKeyConstants.SESSION_STOP_TOPIC, StringCodec.INSTANCE);
  1165 + }
  1166 +
  1167 + /**
  1168 + * 正文末尾追加提示
  1169 + */
  1170 + static Ai1LlmChatTool.ChatText withNotice(Ai1LlmChatTool.ChatText chatText, String notice) {
  1171 + return new Ai1LlmChatTool.ChatText(StrUtil.nullToEmpty(chatText.getContent()) + NOTICE_PREFIX + notice, chatText.getThinking());
  1172 + }
  1173 +
  1174 + /**
  1175 + * 解析正文末尾已回填的提示;没有时返回空
  1176 + */
  1177 + static String getSavedNotice(String content) {
  1178 + for (String notice : List.of(TERMINATED_NOTICE, FAILED_NOTICE, INTERRUPTED_MESSAGE)) {
  1179 + if (StrUtil.endWith(content, NOTICE_PREFIX + notice)) {
  1180 + return notice;
  1181 + }
  1182 + }
  1183 + return null;
  1184 + }
  1185 +
833 1186 /**
834 1187 * 解析结果流条目编号(格式为 毫秒时间戳-序号);为空或非法时从头读取
835 1188 */
... ... @@ -844,4 +1197,94 @@ public class Ai1SessionStreamTool implements SmartLifecycle {
844 1197 return new StreamMessageId(Long.parseLong(parts.get(0)), Long.parseLong(parts.get(1)));
845 1198 }
846 1199  
  1200 + /**
  1201 + * 生成的运行上下文
  1202 + *
  1203 + * 停止与结束互斥,结束后迟到的停止不再中断线程
  1204 + */
  1205 + static class RunContext {
  1206 +
  1207 + /**
  1208 + * 生成线程,即 worker 线程
  1209 + */
  1210 + private final Thread thread;
  1211 + /**
  1212 + * Shell 进程跟踪器
  1213 + */
  1214 + @Getter
  1215 + private final Ai1SkillShellTracker shellTracker;
  1216 + /**
  1217 + * 是否已结束
  1218 + */
  1219 + private boolean finished;
  1220 + /**
  1221 + * 停止提示;为空表示未停止
  1222 + */
  1223 + private String stopNotice;
  1224 +
  1225 + /**
  1226 + * 未结束的 Shell 进程回收线程,与运行上下文的生命周期无关,由节点停止时统一等待
  1227 + */
  1228 + private final Set<Thread> reapers;
  1229 +
  1230 + RunContext(Thread thread, Ai1SkillShellTracker shellTracker, Set<Thread> reapers) {
  1231 + this.thread = thread;
  1232 + this.shellTracker = shellTracker;
  1233 + this.reapers = reapers;
  1234 + }
  1235 +
  1236 + /**
  1237 + * 停止生成:中断生成线程,并由独立线程回收在册的后台进程;已结束、已停止时忽略
  1238 + *
  1239 + * 单个进程回收最多等待 5s:不在调用方执行,避免阻塞停止检查线程或其他上下文的停止;
  1240 + * 不依赖 worker 响应中断,模型或工具库吞掉中断时进程也会被回收;
  1241 + * 回收可能晚于生成结束,结束时 worker 拿不到已被回收线程认领的进程,因此回收线程登记到 reapers,不随上下文移除
  1242 + *
  1243 + * @param notice 停止提示
  1244 + */
  1245 + synchronized void cancel(String notice) {
  1246 + if (finished || stopNotice != null) {
  1247 + return;
  1248 + }
  1249 + stopNotice = notice;
  1250 + thread.interrupt();
  1251 + Thread reaper = ThreadUtil.newThread(() -> {
  1252 + try {
  1253 + shellTracker.cancel();
  1254 + } finally {
  1255 + reapers.remove(Thread.currentThread());
  1256 + }
  1257 + }, REAPER_THREAD_NAME_PREFIX + thread.getName(), true);
  1258 + // 先登记再启动,结束时必能移除自身
  1259 + reapers.add(reaper);
  1260 + try {
  1261 + reaper.start();
  1262 + } catch (Throwable e) {
  1263 + reapers.remove(reaper);
  1264 + throw e;
  1265 + }
  1266 + }
  1267 +
  1268 + /**
  1269 + * 校验生成仍在进行,已停止或已结束时中止后续步骤
  1270 + */
  1271 + synchronized void checkActive() {
  1272 + if (finished || stopNotice != null) {
  1273 + throw new CancellationException(StrUtil.blankToDefault(stopNotice, "生成已结束"));
  1274 + }
  1275 + }
  1276 +
  1277 + /**
  1278 + * 结束生成,并清除 worker 线程的中断标志;需在 worker 线程调用
  1279 + *
  1280 + * @return 停止提示;未停止时为空
  1281 + */
  1282 + synchronized String finish() {
  1283 + finished = true;
  1284 + Thread.interrupted();
  1285 + return stopNotice;
  1286 + }
  1287 +
  1288 + }
  1289 +
847 1290 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTracker.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.skill;
  2 +
  3 +import lombok.extern.slf4j.Slf4j;
  4 +import org.springaicommunity.agent.common.exec.ExecBackend;
  5 +import org.springaicommunity.agent.common.exec.ExecHandle;
  6 +import org.springaicommunity.agent.common.exec.ExecResult;
  7 +import org.springaicommunity.agent.common.exec.ExecSpec;
  8 +
  9 +import java.util.Set;
  10 +import java.util.concurrent.ConcurrentHashMap;
  11 +
  12 +/**
  13 + * AI1 SKILL Shell 进程跟踪器
  14 + *
  15 + * 登记本次生成启动的后台进程,生成结束或终止时回收
  16 + *
  17 + * @author 芋道源码
  18 + */
  19 +@Slf4j
  20 +public class Ai1SkillShellTracker implements ExecBackend {
  21 +
  22 + private final ExecBackend delegate;
  23 + /**
  24 + * 在册的后台进程
  25 + */
  26 + private final Set<ExecHandle> handles = ConcurrentHashMap.newKeySet();
  27 + /**
  28 + * 是否已终止
  29 + */
  30 + private volatile boolean cancelled = false;
  31 +
  32 + public Ai1SkillShellTracker(ExecBackend delegate) {
  33 + this.delegate = delegate;
  34 + }
  35 +
  36 + @Override
  37 + public ExecResult run(ExecSpec spec) {
  38 + return delegate.run(spec);
  39 + }
  40 +
  41 + @Override
  42 + public ExecHandle start(ExecSpec spec) {
  43 + if (cancelled) {
  44 + throw new IllegalStateException("生成已终止");
  45 + }
  46 + ExecHandle handle = delegate.start(spec);
  47 + handles.add(handle);
  48 + // 登记后才被终止时,立即回收;先置标志再遍历的 cancel 与此处必有一方看到该进程
  49 + if (cancelled && handles.remove(handle)) {
  50 + handle.kill();
  51 + }
  52 + return handle;
  53 + }
  54 +
  55 + /**
  56 + * 终止并回收全部在册的后台进程(幂等,可并发调用)
  57 + *
  58 + * 按「移出集合」认领进程,每个进程只回收一次
  59 + */
  60 + public void cancel() {
  61 + cancelled = true;
  62 + for (ExecHandle handle : handles) {
  63 + if (!handles.remove(handle)) {
  64 + continue;
  65 + }
  66 + try {
  67 + handle.kill();
  68 + } catch (Exception e) {
  69 + log.warn("[cancel][Shell 进程({}) 回收失败]", handle.id(), e);
  70 + }
  71 + }
  72 + }
  73 +
  74 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillToolFactory.java
... ... @@ -18,6 +18,7 @@ import jakarta.annotation.Resource;
18 18 import lombok.AllArgsConstructor;
19 19 import lombok.Data;
20 20 import lombok.extern.slf4j.Slf4j;
  21 +import org.springaicommunity.agent.exec.LocalExecBackend;
21 22 import org.springaicommunity.agent.tools.*;
22 23 import org.springframework.ai.tool.ToolCallback;
23 24 import org.springframework.stereotype.Component;
... ... @@ -87,10 +88,11 @@ public class Ai1SkillToolFactory {
87 88 /**
88 89 * 构建 Agent 绑定 SKILL 所需的全部工具:SkillsTool(注入 SKILL.md 内容)+ 执行类工具(限定于 Agent 物化根目录)
89 90 *
90   - * @param agent Agent
  91 + * @param agent Agent
  92 + * @param tracker Shell 进程跟踪器,Shell 命令经其执行,生成终止或结束时统一回收
91 93 * @return 工具列表;未绑定有效 SKILL 时返回空列表
92 94 */
93   - public List<Object> buildTools(Ai1AgentDO agent) {
  95 + public List<Object> buildTools(Ai1AgentDO agent, Ai1SkillShellTracker tracker) {
94 96 // 1. 同步物化快照;未绑定有效 SKILL 时,不提供任何工具
95 97 SkillSnapshot snapshot = sync(agent);
96 98 if (snapshot == null) {
... ... @@ -105,7 +107,7 @@ public class Ai1SkillToolFactory {
105 107  
106 108 // 3. 执行类工具:ShellTools(bash)、FileSystemTools(Read/Write/Edit)、Glob、Grep、ListDirectory
107 109 Path root = snapshot.getRoot();
108   - tools.add(ShellTools.builder().workingDirectory(root).build());
  110 + tools.add(ShellTools.builder().execBackend(tracker).build());
109 111 tools.add(FileSystemTools.builder().allowedDirectory(root).build());
110 112 tools.add(GlobTool.builder().workingDirectory(root).allowedDirectory(root).build());
111 113 tools.add(GrepTool.builder().workingDirectory(root).allowedDirectory(root).build());
... ... @@ -114,6 +116,16 @@ public class Ai1SkillToolFactory {
114 116 }
115 117  
116 118 /**
  119 + * 创建 Agent 的 Shell 进程跟踪器:工作目录限定于 Agent 物化根目录
  120 + *
  121 + * @param agentId Agent 编号
  122 + * @return Shell 进程跟踪器
  123 + */
  124 + public Ai1SkillShellTracker newShellTracker(Long agentId) {
  125 + return new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(buildAgentRoot(agentId)).build());
  126 + }
  127 +
  128 + /**
117 129 * 清理 Agent 物化沙箱:摘除快照缓存并删除本地目录(Agent 删除时调用)
118 130 *
119 131 * @param agentId Agent 编号
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/job/Ai1MemoryArchiveJob.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.job;
  2 +
  3 +import cn.iocoder.yudao.framework.quartz.core.handler.JobHandler;
  4 +import cn.iocoder.yudao.framework.tenant.core.job.TenantJob;
  5 +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryArchiveTool;
  6 +import jakarta.annotation.Resource;
  7 +import org.springframework.stereotype.Component;
  8 +
  9 +/**
  10 + * AI1 长期记忆归档 Job
  11 + *
  12 + * @author 芋道源码
  13 + */
  14 +@Component
  15 +public class Ai1MemoryArchiveJob implements JobHandler {
  16 +
  17 + @Resource
  18 + private Ai1MemoryArchiveTool memoryArchiveTool;
  19 +
  20 + @Override
  21 + @TenantJob
  22 + public String execute(String param) {
  23 + int count = memoryArchiveTool.archive();
  24 + return String.format("AI1 长期记忆归档:归档 %s 个用户", count);
  25 + }
  26 +
  27 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/agent/Ai1AgentService.java
... ... @@ -81,9 +81,9 @@ public interface Ai1AgentService {
81 81 Ai1AgentDO validateAgentEnabled(Long id);
82 82  
83 83 /**
84   - * 获得指定状态的 Agent 列表,用于会话页选择 Agent
  84 + * 获得指定状态的 Agent 列表,按编号升序,用于选择 Agent
85 85 *
86   - * @param status 状态
  86 + * @param status 状态;为空时返回全部状态
87 87 * @return Agent 列表
88 88 */
89 89 List<Ai1AgentDO> getAgentListByStatus(Integer status);
... ...
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;
10 10 import cn.iocoder.yudao.module.ai1.dal.mysql.agent.Ai1AgentMapper;
11 11 import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
12 12 import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory;
  13 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
13 14 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
14 15 import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService;
15 16 import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO;
... ... @@ -45,6 +46,9 @@ public class Ai1AgentServiceImpl implements Ai1AgentService {
45 46 @Resource
46 47 @Lazy // 延迟加载,避免循环依赖
47 48 private Ai1SessionService sessionService;
  49 + @Resource
  50 + @Lazy // 延迟加载,避免循环依赖
  51 + private Ai1MemoryService memoryService;
48 52  
49 53 @Resource
50 54 private Ai1SkillToolFactory skillToolFactory;
... ... @@ -82,9 +86,10 @@ public class Ai1AgentServiceImpl implements Ai1AgentService {
82 86 // 1. 校验存在
83 87 ids.forEach(this::validateAgentExists);
84 88  
85   - // 2. 删除 Agent,同一事务内级联删除会话与消息
  89 + // 2. 删除 Agent,同一事务内级联删除会话、消息与长期记忆
86 90 agentMapper.deleteByIds(ids);
87 91 sessionService.deleteSessionListByAgentIds(ids);
  92 + memoryService.deleteMemoryListByAgentIds(ids);
88 93  
89 94 // 3. 清理 SKILL 物化沙箱
90 95 ids.forEach(skillToolFactory::evict);
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryService.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.service.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  4 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO;
  5 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO;
  6 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  7 +import jakarta.validation.Valid;
  8 +
  9 +import java.util.Collection;
  10 +import java.util.List;
  11 +
  12 +/**
  13 + * AI1 长期记忆 Service 接口
  14 + *
  15 + * @author 芋道源码
  16 + */
  17 +public interface Ai1MemoryService {
  18 +
  19 + /**
  20 + * 修改长期记忆正文
  21 + *
  22 + * @param updateReqVO 修改信息
  23 + */
  24 + void updateMemory(@Valid Ai1MemoryUpdateReqVO updateReqVO);
  25 +
  26 + /**
  27 + * 删除长期记忆(逻辑删除);之后的归档会根据历史消息生成新的记录
  28 + *
  29 + * @param id 编号
  30 + */
  31 + void deleteMemory(Long id);
  32 +
  33 + /**
  34 + * 批量删除长期记忆
  35 + *
  36 + * @param ids 编号列表
  37 + */
  38 + void deleteMemoryListByIds(List<Long> ids);
  39 +
  40 + /**
  41 + * 删除 Agent 下的全部长期记忆(Agent 删除时调用)
  42 + *
  43 + * @param agentIds Agent 编号集合
  44 + */
  45 + void deleteMemoryListByAgentIds(Collection<Long> agentIds);
  46 +
  47 + /**
  48 + * 获得长期记忆
  49 + *
  50 + * @param id 编号
  51 + * @return 长期记忆
  52 + */
  53 + Ai1MemoryDO getMemory(Long id);
  54 +
  55 + /**
  56 + * 获得用户在 Agent 下的长期记忆
  57 + *
  58 + * @param agentId Agent 编号
  59 + * @param userId 用户编号
  60 + * @return 长期记忆
  61 + */
  62 + Ai1MemoryDO getMemoryByAgentIdAndUserId(Long agentId, Long userId);
  63 +
  64 + /**
  65 + * 创建归档的长期记忆;写入前校验,Agent 已删除,或同一用户已有未删除的记忆时不创建
  66 + *
  67 + * @param memory 长期记忆
  68 + * @return 是否创建成功
  69 + */
  70 + boolean createMemoryIfAgentExists(Ai1MemoryDO memory);
  71 +
  72 + /**
  73 + * 更新归档的长期记忆;写入前校验,读取后被人工修改或删除时不更新
  74 + *
  75 + * @param old 归档读取时的记忆
  76 + * @param content 新的记忆正文
  77 + * @param lastMessageId 新的水位线
  78 + * @return 是否更新成功
  79 + */
  80 + boolean updateMemoryIfUnchanged(Ai1MemoryDO old, String content, Long lastMessageId);
  81 +
  82 + /**
  83 + * 获得长期记忆分页
  84 + *
  85 + * @param pageReqVO 分页查询
  86 + * @return 长期记忆分页
  87 + */
  88 + PageResult<Ai1MemoryDO> getMemoryPage(Ai1MemoryPageReqVO pageReqVO);
  89 +
  90 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImpl.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.service.memory;
  2 +
  3 +import cn.hutool.core.collection.CollUtil;
  4 +import cn.hutool.core.util.ObjUtil;
  5 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  6 +import cn.iocoder.yudao.framework.common.util.object.BeanUtils;
  7 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO;
  8 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO;
  9 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  10 +import cn.iocoder.yudao.module.ai1.dal.mysql.memory.Ai1MemoryMapper;
  11 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
  12 +import jakarta.annotation.Resource;
  13 +import org.springframework.context.annotation.Lazy;
  14 +import org.springframework.stereotype.Service;
  15 +import org.springframework.validation.annotation.Validated;
  16 +
  17 +import java.util.Collection;
  18 +import java.util.Collections;
  19 +import java.util.List;
  20 +
  21 +import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
  22 +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MEMORY_NOT_EXISTS;
  23 +
  24 +/**
  25 + * AI1 长期记忆 Service 实现类
  26 + *
  27 + * @author 芋道源码
  28 + */
  29 +@Service
  30 +@Validated
  31 +public class Ai1MemoryServiceImpl implements Ai1MemoryService {
  32 +
  33 + @Resource
  34 + private Ai1MemoryMapper memoryMapper;
  35 +
  36 + @Resource
  37 + @Lazy // 延迟加载,避免循环依赖
  38 + private Ai1AgentService agentService;
  39 +
  40 + @Override
  41 + public void updateMemory(Ai1MemoryUpdateReqVO updateReqVO) {
  42 + // 1. 校验存在
  43 + validateMemoryExists(updateReqVO.getId());
  44 +
  45 + // 2. 更新正文;水位线不变,之后的归档写入前校验到正文变化时放弃
  46 + Ai1MemoryDO updateObj = BeanUtils.toBean(updateReqVO, Ai1MemoryDO.class);
  47 + memoryMapper.updateById(updateObj);
  48 + }
  49 +
  50 + @Override
  51 + public void deleteMemory(Long id) {
  52 + deleteMemoryListByIds(Collections.singletonList(id));
  53 + }
  54 +
  55 + @Override
  56 + public void deleteMemoryListByIds(List<Long> ids) {
  57 + // 1. 校验存在
  58 + if (CollUtil.isEmpty(ids)) {
  59 + return;
  60 + }
  61 + ids.forEach(this::validateMemoryExists);
  62 +
  63 + // 2. 逻辑删除,保留历史记录;进行中的归档写入会因记录已删除而放弃
  64 + memoryMapper.deleteByIds(ids);
  65 + }
  66 +
  67 + @Override
  68 + public void deleteMemoryListByAgentIds(Collection<Long> agentIds) {
  69 + if (CollUtil.isEmpty(agentIds)) {
  70 + return;
  71 + }
  72 + memoryMapper.deleteBatch(Ai1MemoryDO::getAgentId, agentIds);
  73 + }
  74 +
  75 + private void validateMemoryExists(Long id) {
  76 + if (memoryMapper.selectById(id) == null) {
  77 + throw exception(MEMORY_NOT_EXISTS);
  78 + }
  79 + }
  80 +
  81 + @Override
  82 + public Ai1MemoryDO getMemory(Long id) {
  83 + return memoryMapper.selectById(id);
  84 + }
  85 +
  86 + @Override
  87 + public Ai1MemoryDO getMemoryByAgentIdAndUserId(Long agentId, Long userId) {
  88 + return memoryMapper.selectByAgentIdAndUserId(agentId, userId);
  89 + }
  90 +
  91 + @Override
  92 + public boolean createMemoryIfAgentExists(Ai1MemoryDO memory) {
  93 + // 1. 前置校验 Agent 存在,且该用户没有未删除的记忆;同租户归档由归档锁串行,极端并发下的重复记录由按用户查询容忍
  94 + if (agentService.getAgent(memory.getAgentId()) == null
  95 + || memoryMapper.selectByAgentIdAndUserId(memory.getAgentId(), memory.getUserId()) != null) {
  96 + return false;
  97 + }
  98 +
  99 + // 2. 插入
  100 + memoryMapper.insert(memory);
  101 + return true;
  102 + }
  103 +
  104 + @Override
  105 + public boolean updateMemoryIfUnchanged(Ai1MemoryDO old, String content, Long lastMessageId) {
  106 + // 1. 校验记忆未删除,且水位线与正文都未变化
  107 + Ai1MemoryDO memory = memoryMapper.selectById(old.getId());
  108 + if (memory == null || ObjUtil.notEqual(memory.getLastMessageId(), old.getLastMessageId())
  109 + || ObjUtil.notEqual(memory.getContent(), old.getContent())) {
  110 + return false;
  111 + }
  112 +
  113 + // 2. 更新正文并推进水位线;校验之后的人工修改不再拦截,已删除的记录不会被更新
  114 + return memoryMapper.updateById(new Ai1MemoryDO().setId(old.getId())
  115 + .setContent(content).setLastMessageId(lastMessageId)) > 0;
  116 + }
  117 +
  118 + @Override
  119 + public PageResult<Ai1MemoryDO> getMemoryPage(Ai1MemoryPageReqVO pageReqVO) {
  120 + return memoryMapper.selectPage(pageReqVO);
  121 + }
  122 +
  123 +}
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageService.java
1 1 package cn.iocoder.yudao.module.ai1.service.session;
2 2  
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  4 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO;
3 5 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO;
4 6 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
5 7 import jakarta.validation.Valid;
... ... @@ -24,6 +26,14 @@ public interface Ai1SessionMessageService {
24 26 List<Ai1SessionMessageDO> getSessionMessageListBySessionId(Long sessionId);
25 27  
26 28 /**
  29 + * 获得会话消息分页,按编号倒序
  30 + *
  31 + * @param pageReqVO 分页查询
  32 + * @return 消息分页
  33 + */
  34 + PageResult<Ai1SessionMessageDO> getSessionMessagePage(Ai1SessionMessagePageReqVO pageReqVO);
  35 +
  36 + /**
27 37 * 获得会话中指定消息之前最近的若干条消息,按编号倒序
28 38 *
29 39 * @param sessionId 会话编号
... ... @@ -34,6 +44,37 @@ public interface Ai1SessionMessageService {
34 44 List<Ai1SessionMessageDO> getSessionMessageListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit);
35 45  
36 46 /**
  47 + * 获得会话中生成中的助手消息编号列表
  48 + *
  49 + * @param sessionId 会话编号
  50 + * @return 消息编号列表
  51 + */
  52 + List<Long> getGeneratingAssistantMessageIdList(Long sessionId);
  53 +
  54 + /**
  55 + * 获得多个会话中,编号大于 id 的第一条指定状态的消息
  56 + *
  57 + * @param sessionIds 会话编号集合
  58 + * @param status 状态
  59 + * @param id 消息编号(不含)
  60 + * @return 消息
  61 + */
  62 + Ai1SessionMessageDO getFirstSessionMessageBySessionIdsAndStatusAndIdGt(Collection<Long> sessionIds, Integer status, Long id);
  63 +
  64 + /**
  65 + * 获得多个会话中,编号在 (minId, maxId) 之间的指定状态消息,按编号升序
  66 + *
  67 + * @param sessionIds 会话编号集合
  68 + * @param status 状态
  69 + * @param minId 最小编号(不含)
  70 + * @param maxId 最大编号(不含),为空时不限
  71 + * @param limit 数量上限
  72 + * @return 消息列表
  73 + */
  74 + List<Ai1SessionMessageDO> getSessionMessageListBySessionIdsAndStatusAndIdBetween(Collection<Long> sessionIds, Integer status,
  75 + Long minId, Long maxId, Integer limit);
  76 +
  77 + /**
37 78 * 发送消息(SSE 流式):校验会话归属 → 落库用户消息与助手占位 → 投递生成任务 → 打开 SSE 连接
38 79 *
39 80 * 校验、落库失败时不抛出异常,而是返回只包含 error 事件的 SSE 连接
... ... @@ -55,6 +96,14 @@ public interface Ai1SessionMessageService {
55 96 SseEmitter resumeSessionMessageStream(Long userId, Long messageId, String lastEventId);
56 97  
57 98 /**
  99 + * 停止【我的】会话生成:校验会话归属后,终止会话中生成中的助手消息,跨节点生效
  100 + *
  101 + * @param userId 用户编号
  102 + * @param sessionId 会话编号
  103 + */
  104 + void stopSessionMessageMy(Long userId, Long sessionId);
  105 +
  106 + /**
58 107 * 获得消息
59 108 *
60 109 * @param id 编号
... ...
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;
4 4 import cn.hutool.core.util.StrUtil;
5 5 import cn.hutool.extra.spring.SpringUtil;
6 6 import cn.iocoder.yudao.framework.common.exception.ServiceException;
  7 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
7 8 import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
  9 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO;
8 10 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO;
9 11 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
10 12 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
11 13 import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMessageMapper;
  14 +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO;
12 15 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
13 16 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
14 17 import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
... ... @@ -21,6 +24,7 @@ import org.springframework.validation.annotation.Validated;
21 24 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
22 25  
23 26 import java.util.Collection;
  27 +import java.util.Collections;
24 28 import java.util.List;
25 29  
26 30 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
... ... @@ -49,25 +53,57 @@ public class Ai1SessionMessageServiceImpl implements Ai1SessionMessageService {
49 53 @Resource
50 54 private Ai1SessionStreamTool sessionStreamTool;
51 55  
  56 + @Resource
  57 + private Ai1MemoryLockRedisDAO memoryLockRedisDAO;
  58 +
52 59 @Override
53 60 public List<Ai1SessionMessageDO> getSessionMessageListBySessionId(Long sessionId) {
54 61 return sessionMessageMapper.selectListBySessionId(sessionId);
55 62 }
56 63  
57 64 @Override
  65 + public PageResult<Ai1SessionMessageDO> getSessionMessagePage(Ai1SessionMessagePageReqVO pageReqVO) {
  66 + return sessionMessageMapper.selectPage(pageReqVO);
  67 + }
  68 +
  69 + @Override
58 70 public List<Ai1SessionMessageDO> getSessionMessageListBySessionIdAndIdLessThan(Long sessionId, Long id, Integer limit) {
59 71 return sessionMessageMapper.selectListBySessionIdAndIdLessThan(sessionId, id, limit);
60 72 }
61 73  
62 74 @Override
  75 + public List<Long> getGeneratingAssistantMessageIdList(Long sessionId) {
  76 + return sessionMessageMapper.selectIdListBySessionIdAndRoleAndStatus(sessionId,
  77 + Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  78 + }
  79 +
  80 + @Override
  81 + public Ai1SessionMessageDO getFirstSessionMessageBySessionIdsAndStatusAndIdGt(Collection<Long> sessionIds, Integer status, Long id) {
  82 + if (CollUtil.isEmpty(sessionIds)) {
  83 + return null;
  84 + }
  85 + return sessionMessageMapper.selectFirstBySessionIdsAndStatusAndIdGt(sessionIds, status, id);
  86 + }
  87 +
  88 + @Override
  89 + public List<Ai1SessionMessageDO> getSessionMessageListBySessionIdsAndStatusAndIdBetween(Collection<Long> sessionIds, Integer status,
  90 + Long minId, Long maxId, Integer limit) {
  91 + if (CollUtil.isEmpty(sessionIds)) {
  92 + return Collections.emptyList();
  93 + }
  94 + return sessionMessageMapper.selectListBySessionIdsAndStatusAndIdBetween(sessionIds, status, minId, maxId, limit);
  95 + }
  96 +
  97 + @Override
63 98 public SseEmitter sendSessionMessageStream(Long userId, Ai1SessionMessageSendReqVO sendReqVO) {
64 99 try {
65 100 // 1. 校验会话归属、Agent 存在且已开启
66 101 Ai1SessionDO session = sessionService.validateSessionMy(userId, sendReqVO.getSessionId());
67 102 agentService.validateAgentEnabled(session.getAgentId());
68 103  
69   - // 2. 同一事务内落库用户消息与助手占位
70   - Long messageId = getSelf().createRoundMessages(session, sendReqVO.getContent());
  104 + // 2. 同一事务内落库用户消息与助手占位;持有记忆快照锁至事务提交
  105 + Long messageId = memoryLockRedisDAO.lockSnapshot(session.getAgentId(), session.getUserId(),
  106 + () -> getSelf().createRoundMessages(session, sendReqVO.getContent()));
71 107  
72 108 // 3. 投递生成任务,并打开 SSE 连接
73 109 try {
... ... @@ -120,6 +156,14 @@ public class Ai1SessionMessageServiceImpl implements Ai1SessionMessageService {
120 156 }
121 157 }
122 158  
  159 + @Override
  160 + public void stopSessionMessageMy(Long userId, Long sessionId) {
  161 + // 1. 校验会话属于当前用户
  162 + sessionService.validateSessionMy(userId, sessionId);
  163 + // 2. 取消会话中生成中的助手消息
  164 + sessionStreamTool.cancel(sessionId);
  165 + }
  166 +
123 167 /**
124 168 * 开启一轮会话:首条消息自动生成标题 → 落库用户消息(完成)与助手占位(生成中)→ 刷新会话活跃时间
125 169 *
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionService.java
... ... @@ -54,6 +54,14 @@ public interface Ai1SessionService {
54 54 Ai1SessionDO validateSessionMy(Long userId, Long id);
55 55  
56 56 /**
  57 + * 获得会话
  58 + *
  59 + * @param id 编号
  60 + * @return 会话
  61 + */
  62 + Ai1SessionDO getSession(Long id);
  63 +
  64 + /**
57 65 * 获得指定用户在指定 Agent 下的会话列表,按编号倒序
58 66 *
59 67 * @param agentId Agent 编号
... ... @@ -63,6 +71,14 @@ public interface Ai1SessionService {
63 71 List<Ai1SessionDO> getSessionListByAgentIdAndUserId(Long agentId, Long userId);
64 72  
65 73 /**
  74 + * 获得 Agent 下有会话的用户编号列表
  75 + *
  76 + * @param agentId Agent 编号
  77 + * @return 用户编号列表
  78 + */
  79 + List<Long> getUserIdListByAgentId(Long agentId);
  80 +
  81 + /**
66 82 * 获得会话分页(管理端查看 Agent 的全部会话)
67 83 *
68 84 * @param pageReqVO 分页查询
... ...
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 {
84 84 }
85 85  
86 86 @Override
  87 + public Ai1SessionDO getSession(Long id) {
  88 + return sessionMapper.selectById(id);
  89 + }
  90 +
  91 + @Override
87 92 public List<Ai1SessionDO> getSessionListByAgentIdAndUserId(Long agentId, Long userId) {
88 93 return sessionMapper.selectListByAgentIdAndUserId(agentId, userId);
89 94 }
90 95  
91 96 @Override
  97 + public List<Long> getUserIdListByAgentId(Long agentId) {
  98 + return sessionMapper.selectUserIdListByAgentId(agentId);
  99 + }
  100 +
  101 + @Override
92 102 public PageResult<Ai1SessionDO> getSessionPage(Ai1SessionPageReqVO pageReqVO) {
93 103 return sessionMapper.selectPage(pageReqVO);
94 104 }
... ...
yudao-module-ai1-springai/src/main/java/cn/iocoder/yudao/module/ai1/util/Ai1Utils.java
... ... @@ -87,4 +87,56 @@ public class Ai1Utils {
87 87 }
88 88 }
89 89  
  90 + /**
  91 + * 规范化模型生成的长期记忆 Markdown:修正标题、列表标记,并保证维度标题与内容不在同一行
  92 + *
  93 + * 只按长期记忆的固定维度标题拆行,不拆开 C#、#123、URL 锚点,不把中文破折号视为列表标记
  94 + *
  95 + * @param content 模型输出的记忆
  96 + * @param sections 记忆维度标题
  97 + * @return 规范化后的记忆;为空时原样返回
  98 + */
  99 + public static String normalizeMemoryMarkdown(String content, Collection<String> sections) {
  100 + if (StrUtil.isBlank(content)) {
  101 + return content;
  102 + }
  103 + String text = content.replace("\r\n", "\n").replace("\r", "\n").trim();
  104 + // 1. 行内出现的维度标题另起一行:…文字## 用户画像 → …文字\n## 用户画像
  105 + text = text.replaceAll("([^#\\n])(#{1,6}[ \\t]*(?:" + String.join("|", sections) + "))", "$1\n$2");
  106 + // 2. 标题标记后补空格:#标题 → # 标题(# 后为非 #、非空白,避免把 ## 拆成 # #)
  107 + text = text.replaceAll("(?m)^(#{1,6})(?=[^#\\s])", "$1 ");
  108 + // 3. 维度标题与紧随其后的内容拆行:## 偏好习惯-xxx → ## 偏好习惯\n-xxx
  109 + for (String section : sections) {
  110 + text = text.replaceAll("(?m)(#{1,6} " + section + ")[ \\t]*(?=\\S)", "$1\n");
  111 + }
  112 + // 4. 行内列表项另起一行,仅当前一字符为非 ASCII:…口吻-开场 → …口吻\n- 开场(不把中文破折号 —、– 视为列表标记)
  113 + text = text.replaceAll("([^\\x00-\\x7F])[-*+\\u2022][ \\t]*(?=\\S)", "$1\n- ");
  114 + // 5. 行首列表标记后补空格:-xxx → - xxx(同样不处理中文破折号)
  115 + text = text.replaceAll("(?m)^([-*+\\u2022])(?=\\S)", "$1 ");
  116 + // 6. 不同维度之间空一行
  117 + text = text.replaceAll("(?m)^(#{1,6} )", "\n$1");
  118 + // 7. 合并多余空行
  119 + text = text.replaceAll("\\n{3,}", "\n\n");
  120 + return text.trim();
  121 + }
  122 +
  123 + /**
  124 + * 去除首尾空白后按字符数截断,不拆开 UTF-16 代理对
  125 + *
  126 + * @param content 内容
  127 + * @param maxLength 最大长度(字符),需大于 0
  128 + * @return 截断后的内容;为空时原样返回
  129 + */
  130 + public static String limitLength(String content, int maxLength) {
  131 + if (StrUtil.isBlank(content)) {
  132 + return content;
  133 + }
  134 + String text = content.trim();
  135 + if (text.length() <= maxLength) {
  136 + return text;
  137 + }
  138 + int end = Character.isHighSurrogate(text.charAt(maxLength - 1)) ? maxLength - 1 : maxLength;
  139 + return text.substring(0, end);
  140 + }
  141 +
90 142 }
... ...
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;
18 18 import org.springframework.ai.tool.function.FunctionToolCallback;
19 19 import reactor.core.publisher.Flux;
20 20  
  21 +import java.util.ArrayList;
  22 +import java.util.HashMap;
21 23 import java.util.List;
  24 +import java.util.Map;
22 25 import java.util.concurrent.atomic.AtomicInteger;
23 26  
24 27 import static org.junit.jupiter.api.Assertions.*;
... ... @@ -89,6 +92,41 @@ public class Ai1LlmChatToolTest extends BaseMockitoUnitTest {
89 92 assertEquals(1, skillCalls.get());
90 93 }
91 94  
  95 + @Test
  96 + public void testChat_whitespaceChunks() {
  97 + // mock chatModel 和 llmModelFactory 的方法:空格、换行分片也是有效增量,只跳过空串
  98 + ChatModel chatModel = mock(ChatModel.class);
  99 + when(chatModel.getDefaultOptions()).thenReturn(ToolCallingChatOptions.builder().build());
  100 + when(chatModel.stream(any(Prompt.class))).thenReturn(Flux.just(
  101 + buildChunk("想", "你好"), buildChunk(" ", " "), buildChunk("\n", "\n"),
  102 + buildChunk("", "世界"), buildChunk(null, "")));
  103 + ChatClient chatClient = ChatClient.create(chatModel);
  104 + when(llmModelFactory.buildChatClient(any(), any())).thenReturn(chatClient);
  105 + // 准备参数
  106 + List<String> thinkingDeltas = new ArrayList<>();
  107 + List<String> contentDeltas = new ArrayList<>();
  108 +
  109 + // 调用
  110 + Ai1LlmChatTool.ChatText result = llmChatTool.chat(new Ai1ModelRespBO(), "", List.of(), "你好",
  111 + List.of(), List.of(), 1L, thinkingDeltas::add, contentDeltas::add);
  112 + // 断言
  113 + assertEquals("你好 \n世界", result.getContent());
  114 + assertEquals("想 \n", result.getThinking());
  115 + assertEquals(List.of("你好", " ", "\n", "世界"), contentDeltas);
  116 + assertEquals(List.of("想", " ", "\n"), thinkingDeltas);
  117 + }
  118 +
  119 + // ========== 测试数据 ==========
  120 +
  121 + private static ChatResponse buildChunk(String reasoning, String content) {
  122 + Map<String, Object> properties = new HashMap<>();
  123 + if (reasoning != null) {
  124 + properties.put("reasoningContent", reasoning);
  125 + }
  126 + return new ChatResponse(List.of(new Generation(AssistantMessage.builder()
  127 + .content(content).properties(properties).build())));
  128 + }
  129 +
92 130 // ========== 测试工具 ==========
93 131  
94 132 public static class AnnotatedTools {
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryArchiveToolTest.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
  4 +import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
  5 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  6 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  7 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  8 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
  9 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
  10 +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO;
  11 +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
  12 +import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
  13 +import cn.iocoder.yudao.module.ai1.harness.session.Ai1SessionStreamTool;
  14 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
  15 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
  16 +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService;
  17 +import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
  18 +import org.junit.jupiter.api.AfterEach;
  19 +import org.junit.jupiter.api.BeforeEach;
  20 +import org.junit.jupiter.api.Test;
  21 +import org.mockito.InOrder;
  22 +import org.mockito.InjectMocks;
  23 +import org.mockito.Mock;
  24 +import org.mockito.Spy;
  25 +import org.redisson.api.RLock;
  26 +import org.redisson.api.RedissonClient;
  27 +import org.springframework.test.util.ReflectionTestUtils;
  28 +
  29 +import java.util.List;
  30 +
  31 +import static org.junit.jupiter.api.Assertions.*;
  32 +import static org.mockito.ArgumentMatchers.*;
  33 +import static org.mockito.Mockito.*;
  34 +
  35 +/**
  36 + * {@link Ai1MemoryArchiveTool} 的单元测试
  37 + *
  38 + * @author 芋道源码
  39 + */
  40 +public class Ai1MemoryArchiveToolTest extends BaseMockitoUnitTest {
  41 +
  42 + private static final Long TENANT_ID = 1L;
  43 + private static final Long AGENT_ID = 1L;
  44 + private static final Long USER_ID = 100L;
  45 + private static final List<Long> SESSION_IDS = List.of(10L, 11L);
  46 +
  47 + @InjectMocks
  48 + private Ai1MemoryArchiveTool memoryArchiveTool;
  49 +
  50 + @Spy
  51 + private Ai1MemoryLockRedisDAO memoryLockRedisDAO;
  52 + @Mock
  53 + private RedissonClient redissonClient;
  54 + @Mock
  55 + private RLock archiveLock;
  56 + @Mock
  57 + private RLock snapshotLock;
  58 + @Spy
  59 + private YudaoAi1Properties ai1Properties = new YudaoAi1Properties();
  60 +
  61 + @Mock
  62 + private Ai1AgentService agentService;
  63 + @Mock
  64 + private Ai1SessionService sessionService;
  65 + @Mock
  66 + private Ai1SessionMessageService sessionMessageService;
  67 + @Mock
  68 + private Ai1MemoryService memoryService;
  69 + @Mock
  70 + private Ai1MemoryTool memoryTool;
  71 + @Mock
  72 + private Ai1SessionStreamTool sessionStreamTool;
  73 +
  74 + @BeforeEach
  75 + public void before() {
  76 + TenantContextHolder.setTenantId(TENANT_ID);
  77 + // mock 归档锁、快照锁
  78 + ReflectionTestUtils.setField(memoryLockRedisDAO, "redissonClient", redissonClient);
  79 + lenient().when(redissonClient.getLock("ai1:memory:archive-lock:" + TENANT_ID)).thenReturn(archiveLock);
  80 + lenient().when(archiveLock.tryLock()).thenReturn(true);
  81 + lenient().when(archiveLock.isHeldByCurrentThread()).thenReturn(true);
  82 + lenient().when(redissonClient.getLock("ai1:memory:snapshot-lock:" + AGENT_ID + ":" + USER_ID)).thenReturn(snapshotLock);
  83 + lenient().when(sessionService.getSessionListByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(List.of(
  84 + new Ai1SessionDO().setId(10L), new Ai1SessionDO().setId(11L)));
  85 + }
  86 +
  87 + @AfterEach
  88 + public void after() {
  89 + TenantContextHolder.clear();
  90 + }
  91 +
  92 + @Test
  93 + public void testArchive_locked() {
  94 + // mock 其他节点归档中
  95 + when(archiveLock.tryLock()).thenReturn(false);
  96 +
  97 + // 调用,并断言
  98 + assertEquals(0, memoryArchiveTool.archive());
  99 + verifyNoInteractions(agentService, memoryTool);
  100 + verify(archiveLock, never()).unlock();
  101 + }
  102 +
  103 + @Test
  104 + public void testArchive_lockNotHeld() {
  105 + // mock 方法:归档结束时锁已不属于当前线程
  106 + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of());
  107 + when(archiveLock.isHeldByCurrentThread()).thenReturn(false);
  108 +
  109 + // 调用,并断言:不释放其他线程持有的锁
  110 + assertEquals(0, memoryArchiveTool.archive());
  111 + verify(archiveLock, never()).unlock();
  112 + }
  113 +
  114 + @Test
  115 + public void testArchive_onlyMemoryEnabled() {
  116 + // mock 方法:只有开启记忆的 Agent 归档
  117 + Ai1AgentDO agent = buildAgent();
  118 + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of(
  119 + agent, buildAgent().setId(2L).setMemoryEnabled(false), buildAgent().setId(3L).setMemoryEnabled(null)));
  120 + when(sessionService.getUserIdListByAgentId(AGENT_ID)).thenReturn(List.of(USER_ID));
  121 + mockNewMessages(0L, null, List.of(buildMessage(12L), buildMessage(13L)));
  122 + when(memoryTool.generate(eq(agent), isNull(), anyList(), eq(2000))).thenReturn("## 用户画像\n- 程序员");
  123 + when(memoryService.createMemoryIfAgentExists(any())).thenReturn(true);
  124 +
  125 + // 调用
  126 + int count = memoryArchiveTool.archive();
  127 + // 断言:首次归档插入,水位线为本批最后一条
  128 + assertEquals(1, count);
  129 + verify(sessionService, never()).getUserIdListByAgentId(2L);
  130 + verify(sessionService, never()).getUserIdListByAgentId(3L);
  131 + verify(memoryService).createMemoryIfAgentExists(argThat(memory -> AGENT_ID.equals(memory.getAgentId())
  132 + && USER_ID.equals(memory.getUserId()) && "## 用户画像\n- 程序员".equals(memory.getContent())
  133 + && Long.valueOf(13L).equals(memory.getLastMessageId())));
  134 + verify(archiveLock).unlock();
  135 + }
  136 +
  137 + @Test
  138 + public void testArchive_userFailedNotAffectOthers() {
  139 + // mock 方法:第一个用户生成失败,第二个用户无新消息
  140 + when(agentService.getAgentListByStatus(CommonStatusEnum.ENABLE.getStatus())).thenReturn(List.of(buildAgent()));
  141 + when(sessionService.getUserIdListByAgentId(AGENT_ID)).thenReturn(List.of(USER_ID, 200L));
  142 + mockNewMessages(0L, null, List.of(buildMessage(12L)));
  143 + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenThrow(new IllegalStateException("模型超时"));
  144 + when(redissonClient.getLock("ai1:memory:snapshot-lock:" + AGENT_ID + ":200")).thenReturn(snapshotLock);
  145 +
  146 + // 调用
  147 + int count = memoryArchiveTool.archive();
  148 + // 断言:不推进水位线,下次重试
  149 + assertEquals(0, count);
  150 + verify(memoryService).getMemoryByAgentIdAndUserId(AGENT_ID, 200L);
  151 + verify(memoryService, never()).createMemoryIfAgentExists(any());
  152 + verify(memoryService, never()).updateMemoryIfUnchanged(any(), any(), any());
  153 + verify(archiveLock).unlock();
  154 + }
  155 +
  156 + @Test
  157 + public void testArchiveOne_safePrefix() {
  158 + // mock 方法:水位线之后有生成中的消息,只读取它之前的已完成消息
  159 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).agentId(AGENT_ID).userId(USER_ID)
  160 + .content("旧记忆").lastMessageId(11L).build();
  161 + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory);
  162 + mockNewMessages(11L, 20L, List.of(buildMessage(12L), buildMessage(13L)));
  163 + when(memoryTool.generate(any(), eq("旧记忆"), anyList(), eq(2000))).thenReturn("新记忆");
  164 + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 13L)).thenReturn(true);
  165 +
  166 + // 调用,并断言
  167 + assertTrue(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  168 + // 断言:孤儿收尾在快照锁外,持有快照锁读取,模型在锁外调用
  169 + InOrder inOrder = inOrder(sessionStreamTool, snapshotLock, sessionMessageService, memoryTool);
  170 + inOrder.verify(sessionStreamTool).failStaleMessage(argThat(message -> Long.valueOf(20L).equals(message.getId())));
  171 + inOrder.verify(snapshotLock).lock();
  172 + inOrder.verify(sessionMessageService).getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS,
  173 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 11L, 20L, Ai1MemoryArchiveTool.BATCH_SIZE);
  174 + inOrder.verify(snapshotLock).unlock();
  175 + inOrder.verify(memoryTool).generate(any(), any(), anyList(), anyInt());
  176 + }
  177 +
  178 + @Test
  179 + public void testArchiveOne_snapshotReadFailed() {
  180 + // mock 方法:持有快照锁读取新消息时失败
  181 + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS,
  182 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 0L, null, Ai1MemoryArchiveTool.BATCH_SIZE))
  183 + .thenThrow(new IllegalStateException("数据库超时"));
  184 +
  185 + // 调用,并断言:异常抛出,快照锁仍释放
  186 + assertThrows(IllegalStateException.class, () -> memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  187 + verify(snapshotLock).lock();
  188 + verify(snapshotLock).unlock();
  189 + verifyNoInteractions(memoryTool);
  190 + }
  191 +
  192 + @Test
  193 + public void testArchiveOne_noNewMessages() {
  194 + // mock 方法
  195 + mockNewMessages(0L, null, List.of());
  196 +
  197 + // 调用,并断言
  198 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  199 + verifyNoInteractions(memoryTool, sessionStreamTool);
  200 + }
  201 +
  202 + @Test
  203 + public void testArchiveOne_staleGeneratingRecovered() {
  204 + // mock 数据:孤儿生成之后有已完成的消息
  205 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).content("旧记忆").lastMessageId(11L).build();
  206 + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory);
  207 + // mock 方法:孤儿生成收尾后,不再有生成中的消息
  208 + Ai1SessionMessageDO generatingMessage = buildMessage(20L);
  209 + when(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(SESSION_IDS,
  210 + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), 11L)).thenReturn(generatingMessage, (Ai1SessionMessageDO) null);
  211 + when(sessionStreamTool.failStaleMessage(generatingMessage)).thenReturn(true);
  212 + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS,
  213 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 11L, null, Ai1MemoryArchiveTool.BATCH_SIZE))
  214 + .thenReturn(List.of(buildMessage(12L), buildMessage(21L)));
  215 + when(memoryTool.generate(any(), eq("旧记忆"), anyList(), eq(2000))).thenReturn("新记忆");
  216 + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 21L)).thenReturn(true);
  217 +
  218 + // 调用,并断言:水位线越过已收尾的孤儿生成
  219 + assertTrue(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  220 + InOrder inOrder = inOrder(sessionStreamTool, snapshotLock);
  221 + inOrder.verify(sessionStreamTool).failStaleMessage(generatingMessage);
  222 + inOrder.verify(snapshotLock).lock();
  223 + }
  224 +
  225 + @Test
  226 + public void testArchiveOne_activeGeneratingBlocks() {
  227 + // mock 方法:生成仍在执行或未超时,安全前缀为空
  228 + mockNewMessages(0L, 12L, List.of());
  229 +
  230 + // 调用,并断言:不收尾,也不调用模型
  231 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  232 + verify(sessionStreamTool).failStaleMessage(argThat(message -> Long.valueOf(12L).equals(message.getId())));
  233 + verifyNoInteractions(memoryTool);
  234 + }
  235 +
  236 + @Test
  237 + public void testArchiveOne_blankMemory() {
  238 + // mock 方法
  239 + mockNewMessages(0L, null, List.of(buildMessage(12L)));
  240 + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn(" ");
  241 +
  242 + // 调用,并断言:模型无输出时不写入,水位线不前进
  243 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  244 + verify(memoryService, never()).createMemoryIfAgentExists(any());
  245 + }
  246 +
  247 + @Test
  248 + public void testArchiveOne_changedDuringArchive() {
  249 + // mock 方法:归档期间记忆被人工修改或删除
  250 + Ai1MemoryDO memory = Ai1MemoryDO.builder().id(1L).content("旧记忆").lastMessageId(11L).build();
  251 + when(memoryService.getMemoryByAgentIdAndUserId(AGENT_ID, USER_ID)).thenReturn(memory);
  252 + mockNewMessages(11L, null, List.of(buildMessage(12L)));
  253 + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn("新记忆");
  254 + when(memoryService.updateMemoryIfUnchanged(memory, "新记忆", 12L)).thenReturn(false);
  255 +
  256 + // 调用,并断言:放弃本次结果,不重新插入
  257 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  258 + verify(memoryService, never()).createMemoryIfAgentExists(any());
  259 + }
  260 +
  261 + @Test
  262 + public void testArchiveOne_insertConflict() {
  263 + // mock 方法:其他归档已插入,或 Agent 已删除
  264 + mockNewMessages(0L, null, List.of(buildMessage(12L)));
  265 + when(memoryTool.generate(any(), any(), anyList(), anyInt())).thenReturn("新记忆");
  266 + when(memoryService.createMemoryIfAgentExists(any())).thenReturn(false);
  267 +
  268 + // 调用,并断言:放弃本次结果
  269 + assertFalse(memoryArchiveTool.archiveOne(buildAgent(), USER_ID, 2000));
  270 + verify(memoryService, never()).updateMemoryIfUnchanged(any(), any(), any());
  271 + }
  272 +
  273 + // ========== mock 方法 ==========
  274 +
  275 + private void mockNewMessages(Long lastMessageId, Long generatingMessageId, List<Ai1SessionMessageDO> messages) {
  276 + when(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(SESSION_IDS,
  277 + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), lastMessageId))
  278 + .thenReturn(generatingMessageId != null ? buildMessage(generatingMessageId) : null);
  279 + when(sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(SESSION_IDS,
  280 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), lastMessageId, generatingMessageId, Ai1MemoryArchiveTool.BATCH_SIZE))
  281 + .thenReturn(messages);
  282 + }
  283 +
  284 + // ========== 测试数据 ==========
  285 +
  286 + private static Ai1AgentDO buildAgent() {
  287 + return new Ai1AgentDO().setId(AGENT_ID).setProviderId(1L).setModelId(2L).setMemoryEnabled(true);
  288 + }
  289 +
  290 + private static Ai1SessionMessageDO buildMessage(Long id) {
  291 + return new Ai1SessionMessageDO().setId(id).setContent("消息" + id);
  292 + }
  293 +
  294 +}
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/memory/Ai1MemoryToolTest.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  4 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  5 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  6 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
  7 +import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
  8 +import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
  9 +import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool;
  10 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
  11 +import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService;
  12 +import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO;
  13 +import org.junit.jupiter.api.Test;
  14 +import org.mockito.ArgumentCaptor;
  15 +import org.mockito.InjectMocks;
  16 +import org.mockito.Mock;
  17 +
  18 +import java.util.List;
  19 +
  20 +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException;
  21 +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MODEL_TYPE_NOT_CHAT;
  22 +import static org.junit.jupiter.api.Assertions.*;
  23 +import static org.mockito.ArgumentMatchers.*;
  24 +import static org.mockito.Mockito.*;
  25 +
  26 +/**
  27 + * {@link Ai1MemoryTool} 的单元测试
  28 + *
  29 + * @author 芋道源码
  30 + */
  31 +public class Ai1MemoryToolTest extends BaseMockitoUnitTest {
  32 +
  33 + @InjectMocks
  34 + private Ai1MemoryTool memoryTool;
  35 +
  36 + @Mock
  37 + private Ai1MemoryService memoryService;
  38 + @Mock
  39 + private Ai1ModelService modelService;
  40 + @Mock
  41 + private Ai1LlmChatTool llmChatTool;
  42 +
  43 + @Test
  44 + public void testLoadMemory() {
  45 + // mock 方法
  46 + when(memoryService.getMemoryByAgentIdAndUserId(1L, 100L)).thenReturn(Ai1MemoryDO.builder().content("## 用户画像\n- 程序员").build());
  47 +
  48 + // 调用,并断言
  49 + assertEquals("## 用户画像\n- 程序员", memoryTool.loadMemory(1L, 100L));
  50 + // 调用,并断言:编号为空时不查询
  51 + assertNull(memoryTool.loadMemory(null, 100L));
  52 + verify(memoryService, never()).getMemoryByAgentIdAndUserId(isNull(), any());
  53 + }
  54 +
  55 + @Test
  56 + public void testMergePrompt() {
  57 + // 调用,并断言:按模板合入记忆
  58 + assertEquals("你是客服\n\n以下是该用户的历史长期记忆(Markdown),请在回答时结合这些信息;"
  59 + + "若与用户当前表述冲突,以当前表述为准:\n## 用户画像\n- 程序员", memoryTool.mergePrompt("你是客服", " ## 用户画像\n- 程序员 \n"));
  60 + // 调用,并断言:记忆为空时原样返回
  61 + assertEquals("你是客服", memoryTool.mergePrompt("你是客服", " "));
  62 + assertEquals("", memoryTool.mergePrompt(null, null));
  63 + }
  64 +
  65 + @Test
  66 + public void testGenerate_modelNotChat() {
  67 + // mock modelService 的方法
  68 + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO()
  69 + .setModelType(Ai1ModelTypeEnum.EMBEDDING.getType()));
  70 +
  71 + // 调用,并断言异常
  72 + assertServiceException(() -> memoryTool.generate(buildAgent(), null, List.of(), 2000), MODEL_TYPE_NOT_CHAT);
  73 + verifyNoInteractions(llmChatTool);
  74 + }
  75 +
  76 + @Test
  77 + public void testGenerate_prompt() {
  78 + // mock 方法
  79 + mockChat("## 用户画像\n- 程序员");
  80 + // 准备参数
  81 + List<Ai1SessionMessageDO> messages = List.of(
  82 + new Ai1SessionMessageDO().setRole(Ai1SessionMessageRoleEnum.USER.getRole()).setContent("我是程序员"),
  83 + new Ai1SessionMessageDO().setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole()).setContent("好的"));
  84 +
  85 + // 调用
  86 + String memory = memoryTool.generate(buildAgent(), " 旧记忆 ", messages, 500);
  87 + // 断言:单轮会话,不带历史、工具、RAG,也不替换会话 Header
  88 + assertEquals("## 用户画像\n- 程序员", memory);
  89 + ArgumentCaptor<String> systemPromptCaptor = ArgumentCaptor.forClass(String.class);
  90 + ArgumentCaptor<String> userPromptCaptor = ArgumentCaptor.forClass(String.class);
  91 + verify(llmChatTool).chat(any(), systemPromptCaptor.capture(), eq(List.of()), userPromptCaptor.capture(),
  92 + eq(List.of()), eq(List.of()), isNull(), isNull(), isNull());
  93 + assertTrue(systemPromptCaptor.getValue().endsWith("总长度不超过 500 字。"));
  94 + assertEquals("【历史记忆】\n旧记忆\n\n【最新对话】\n用户:我是程序员\n助手:好的\n", userPromptCaptor.getValue());
  95 + }
  96 +
  97 + @Test
  98 + public void testGenerate_normalizeSections() {
  99 + // mock 方法:标题与内容压成一行
  100 + mockChat("##用户画像-后端工程师##偏好习惯-简洁回答");
  101 +
  102 + // 调用
  103 + String memory = memoryTool.generate(buildAgent(), null, List.of(), 2000);
  104 + // 断言
  105 + assertEquals("## 用户画像\n- 后端工程师\n\n## 偏好习惯\n- 简洁回答", memory);
  106 + }
  107 +
  108 + @Test
  109 + public void testGenerate_defaultMaxLength() {
  110 + // mock 方法
  111 + mockChat("用户喜欢😀表情");
  112 +
  113 + // 调用:未配置长度上限时使用默认值
  114 + String memory = memoryTool.generate(buildAgent(), null, List.of(), 0);
  115 + // 断言
  116 + assertEquals("用户喜欢😀表情", memory);
  117 + ArgumentCaptor<String> systemPromptCaptor = ArgumentCaptor.forClass(String.class);
  118 + verify(llmChatTool).chat(any(), systemPromptCaptor.capture(), anyList(), anyString(), anyList(), anyList(),
  119 + isNull(), isNull(), isNull());
  120 + assertTrue(systemPromptCaptor.getValue().endsWith("总长度不超过 2000 字。"));
  121 + // 调用,并断言:按长度上限截断
  122 + assertEquals("用户喜欢", memoryTool.generate(buildAgent(), null, List.of(), 5));
  123 + }
  124 +
  125 + @Test
  126 + public void testGenerate_empty() {
  127 + // mock 方法
  128 + mockChat(" ");
  129 +
  130 + // 调用,并断言:模型无输出时返回空,由调用方跳过写入
  131 + assertTrue(memoryTool.generate(buildAgent(), null, List.of(), 2000).isBlank());
  132 + }
  133 +
  134 + // ========== mock 方法 ==========
  135 +
  136 + private void mockChat(String content) {
  137 + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO()
  138 + .setModelType(Ai1ModelTypeEnum.CHAT.getType()));
  139 + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), isNull(), isNull(), isNull()))
  140 + .thenReturn(new Ai1LlmChatTool.ChatText(content, ""));
  141 + }
  142 +
  143 + // ========== 测试数据 ==========
  144 +
  145 + private static Ai1AgentDO buildAgent() {
  146 + return new Ai1AgentDO().setId(1L).setProviderId(1L).setModelId(2L);
  147 + }
  148 +
  149 +}
... ...
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;
3 3 import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
4 4 import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
5 5 import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  6 +import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
6 7 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
7 8 import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
8 9 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
... ... @@ -10,7 +11,9 @@ import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
10 11 import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
11 12 import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool;
12 13 import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory;
  14 +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool;
13 15 import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool;
  16 +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker;
14 17 import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory;
15 18 import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
16 19 import cn.iocoder.yudao.module.ai1.service.knowledge.Ai1KnowledgeBaseService;
... ... @@ -20,29 +23,52 @@ import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionMessageService;
20 23 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
21 24 import org.junit.jupiter.api.BeforeEach;
22 25 import org.junit.jupiter.api.Test;
  26 +import org.junit.jupiter.api.Timeout;
  27 +import org.junit.jupiter.api.condition.EnabledOnOs;
  28 +import org.junit.jupiter.api.condition.OS;
  29 +import org.junit.jupiter.api.io.TempDir;
23 30 import org.junit.jupiter.params.ParameterizedTest;
24 31 import org.junit.jupiter.params.provider.ValueSource;
25 32 import org.mockito.ArgumentCaptor;
26 33 import org.mockito.InjectMocks;
27 34 import org.mockito.Mock;
28 35 import org.mockito.Spy;
  36 +import org.redisson.api.RBucket;
29 37 import org.redisson.api.RLock;
30 38 import org.redisson.api.RStream;
  39 +import org.redisson.api.RTopic;
31 40 import org.redisson.api.RedissonClient;
32 41 import org.redisson.api.stream.StreamAddArgs;
  42 +import org.redisson.api.stream.StreamCreateGroupArgs;
33 43 import org.redisson.api.stream.StreamMessageId;
34 44 import org.redisson.api.stream.StreamRangeArgs;
35 45 import org.redisson.api.stream.StreamReadArgs;
  46 +import org.redisson.api.stream.StreamReadGroupArgs;
  47 +import org.redisson.client.RedisException;
36 48 import org.redisson.client.codec.Codec;
  49 +import org.springaicommunity.agent.common.exec.ExecBackend;
  50 +import org.springaicommunity.agent.common.exec.ExecHandle;
  51 +import org.springaicommunity.agent.common.exec.ExecSpec;
  52 +import org.springaicommunity.agent.exec.LocalExecBackend;
37 53 import org.springframework.test.util.ReflectionTestUtils;
38 54 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
39 55  
40 56 import java.io.IOException;
  57 +import java.nio.file.Files;
  58 +import java.nio.file.Path;
  59 +import java.time.Duration;
  60 +import java.time.LocalDateTime;
41 61 import java.util.ArrayList;
42 62 import java.util.HashMap;
43 63 import java.util.LinkedHashMap;
44 64 import java.util.List;
45 65 import java.util.Map;
  66 +import java.util.Optional;
  67 +import java.util.Set;
  68 +import java.util.concurrent.CancellationException;
  69 +import java.util.concurrent.CountDownLatch;
  70 +import java.util.concurrent.TimeUnit;
  71 +import java.util.concurrent.atomic.AtomicBoolean;
46 72 import java.util.concurrent.atomic.AtomicReference;
47 73 import java.util.function.Consumer;
48 74  
... ... @@ -74,6 +100,10 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
74 100 private RStream<String, String> resultStream;
75 101 @Mock
76 102 private RLock taskLock;
  103 + @Mock
  104 + private RBucket<String> stopBucket;
  105 + @Mock
  106 + private RTopic stopTopic;
77 107 @Spy
78 108 private YudaoAi1Properties ai1Properties = new YudaoAi1Properties();
79 109  
... ... @@ -95,15 +125,24 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
95 125 private Ai1SkillToolFactory skillToolFactory;
96 126 @Mock
97 127 private Ai1RagTool ragTool;
  128 + @Mock
  129 + private Ai1MemoryTool memoryTool;
  130 + @Mock
  131 + private Ai1SkillShellTracker shellTracker;
98 132  
99 133 @BeforeEach
100 134 public void before() {
  135 + // 直接调用 handleTask 的用例按运行中处理,与 start 之后一致
  136 + ReflectionTestUtils.setField(sessionStreamTool, "running", true);
101 137 // mock redissonClient 的方法(部分用例不访问 Redis)
102 138 lenient().doReturn(taskStream).when(redissonClient).getStream(eq("ai1:session:tasks"), any(Codec.class));
103 139 lenient().doReturn(resultStream).when(redissonClient).getStream(eq("ai1:session:result:" + MESSAGE_ID), any(Codec.class));
104 140 lenient().when(redissonClient.getLock("ai1:session:task-lock:" + MESSAGE_ID)).thenReturn(taskLock);
105 141 lenient().when(taskLock.tryLock()).thenReturn(true);
106 142 lenient().when(taskLock.isHeldByThread(anyLong())).thenReturn(true);
  143 + lenient().doReturn(stopBucket).when(redissonClient).getBucket(eq("ai1:session:stop:" + MESSAGE_ID), any(Codec.class));
  144 + lenient().when(redissonClient.getTopic(eq("ai1:session:stop"), any(Codec.class))).thenReturn(stopTopic);
  145 + lenient().when(skillToolFactory.newShellTracker(AGENT_ID)).thenReturn(shellTracker);
107 146 lenient().when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO()
108 147 .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
109 148 .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus()));
... ... @@ -256,11 +295,14 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
256 295  
257 296 // 调用
258 297 sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
259   - // 断言
  298 + // 断言:正文追加固定失败提示,不拼接原始异常
260 299 ArgumentCaptor<Ai1SessionMessageDO> messageCaptor = ArgumentCaptor.forClass(Ai1SessionMessageDO.class);
261 300 verify(sessionMessageService).updateSessionMessage(messageCaptor.capture());
262 301 assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), messageCaptor.getValue().getStatus());
  302 + assertEquals("\n\n> 本条回复生成失败", messageCaptor.getValue().getContent());
263 303 assertEquals(List.of("ping", "error", "done"), captureResultTypes());
  304 + assertEquals(List.of("", "本条回复生成失败", ""), captureResultData());
  305 + verify(shellTracker).cancel();
264 306 }
265 307  
266 308 @Test
... ... @@ -310,7 +352,7 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
310 352 sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), reclaimed);
311 353 // 断言
312 354 verifyNoInteractions(llmChatTool, agentService);
313   - verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段第二段".equals(message.getContent())
  355 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段第二段\n\n> 生成任务已中断,请重新发送".equals(message.getContent())
314 356 && "思考".equals(message.getReasoning())
315 357 && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
316 358 assertEquals(List.of("error", "done"), captureResultTypes());
... ... @@ -326,7 +368,7 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
326 368 sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true);
327 369 // 断言
328 370 verifyNoInteractions(llmChatTool);
329   - verify(sessionMessageService).updateSessionMessage(argThat(message -> "".equals(message.getContent())
  371 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 生成任务已中断,请重新发送".equals(message.getContent())
330 372 && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
331 373 assertEquals(List.of("error", "done"), captureResultTypes());
332 374 }
... ... @@ -437,6 +479,573 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
437 479 verify(taskStream).ack("ai1-session-workers", RECORD_ID);
438 480 }
439 481  
  482 + @Test
  483 + public void testHandleTask_terminatedWithoutTerminal() {
  484 + // mock 方法:人工终止已回填,done 尚未写入
  485 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO()
  486 + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  487 + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent("已生成\n\n> 本条回复已被人工终止"));
  488 + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of());
  489 +
  490 + // 调用
  491 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true);
  492 + // 断言:人工终止不补写 error
  493 + verifyNoInteractions(llmChatTool);
  494 + verify(sessionMessageService, never()).updateSessionMessage(any());
  495 + assertEquals(List.of("done"), captureResultTypes());
  496 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  497 + }
  498 +
  499 + @Test
  500 + public void testHandleTask_generateFailedWithoutTerminal() {
  501 + // mock 方法:生成失败已回填,error 尚未写入
  502 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO()
  503 + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  504 + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent("\n\n> 本条回复生成失败"));
  505 + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of());
  506 +
  507 + // 调用
  508 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields(), true);
  509 + // 断言:按已回填的固定提示补写,不改为中断
  510 + assertEquals(List.of("error", "done"), captureResultTypes());
  511 + assertEquals(List.of("本条回复生成失败", ""), captureResultData());
  512 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  513 + }
  514 +
  515 + @Test
  516 + public void testHandleTask_failedBeforeStart() {
  517 + // mock 方法:投递超时后补偿已标记失败,worker 才开始处理
  518 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO()
  519 + .setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  520 + .setStatus(Ai1SessionMessageStatusEnum.FAILED.getStatus()).setContent(""));
  521 + when(resultStream.rangeReversed(StreamRangeArgs.startId(StreamMessageId.MAX).endId(StreamMessageId.MIN).count(1))).thenReturn(Map.of());
  522 +
  523 + // 调用
  524 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  525 + // 断言:不生成、不改写消息,补齐终态后确认
  526 + verifyNoInteractions(agentService, llmChatTool, mcpToolFactory, skillToolFactory);
  527 + verify(sessionMessageService, never()).updateSessionMessage(any());
  528 + assertEquals(List.of("error", "done"), captureResultTypes());
  529 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  530 + }
  531 +
  532 + @Test
  533 + public void testHandleTask_injectMemory() {
  534 + // mock 方法
  535 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L).setMemoryEnabled(true));
  536 + when(sessionService.getSession(SESSION_ID)).thenReturn(new Ai1SessionDO().setId(SESSION_ID).setUserId(100L));
  537 + when(memoryTool.loadMemory(AGENT_ID, 100L)).thenReturn("## 用户画像\n- 程序员");
  538 + when(memoryTool.mergePrompt("你是 客服 的智能助手。", "## 用户画像\n- 程序员")).thenReturn("合并后的系统指令");
  539 + when(llmChatTool.chat(any(), eq("合并后的系统指令"), anyList(), eq("你好"), anyList(), anyList(), eq(SESSION_ID), any(), any()))
  540 + .thenReturn(new Ai1LlmChatTool.ChatText("您好", ""));
  541 +
  542 + // 调用
  543 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  544 + // 断言
  545 + verify(sessionMessageService).updateSessionMessage(argThat(message ->
  546 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus().equals(message.getStatus())));
  547 + verify(skillToolFactory).buildTools(any(), same(shellTracker));
  548 + }
  549 +
  550 + @Test
  551 + public void testHandleTask_injectMemoryFailed() {
  552 + // mock 方法:记忆查询失败时降级为原系统指令
  553 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L).setMemoryEnabled(true));
  554 + when(sessionService.getSession(SESSION_ID)).thenThrow(new IllegalStateException("数据库不可用"));
  555 + when(llmChatTool.chat(any(), eq("你是 客服 的智能助手。"), anyList(), eq("你好"), anyList(), anyList(), eq(SESSION_ID), any(), any()))
  556 + .thenReturn(new Ai1LlmChatTool.ChatText("您好", ""));
  557 +
  558 + // 调用
  559 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  560 + // 断言
  561 + verifyNoInteractions(memoryTool);
  562 + verify(sessionMessageService).updateSessionMessage(argThat(message ->
  563 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus().equals(message.getStatus())));
  564 + }
  565 +
  566 + @Test
  567 + public void testHandleTask_stoppedWhileQueued() {
  568 + // mock 排队期间已写入停止意图
  569 + when(stopBucket.isExists()).thenReturn(true);
  570 +
  571 + // 调用
  572 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  573 + // 断言:不调用模型,人工终止只写 done,并清除中断标志
  574 + verifyNoInteractions(agentService, llmChatTool);
  575 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 本条回复已被人工终止".equals(message.getContent())
  576 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  577 + assertEquals(List.of("done"), captureResultTypes());
  578 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  579 + assertFalse(Thread.currentThread().isInterrupted());
  580 + }
  581 +
  582 + @Test
  583 + public void testHandleTask_stoppedDuringPreparation() {
  584 + // mock 方法:装配工具期间收到停止通知
  585 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L));
  586 + when(mcpToolFactory.buildTools(any())).thenAnswer(invocation -> {
  587 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID);
  588 + return List.of();
  589 + });
  590 +
  591 + // 调用
  592 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  593 + // 断言:不再调用模型
  594 + verifyNoInteractions(llmChatTool);
  595 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 本条回复已被人工终止".equals(message.getContent())
  596 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  597 + assertEquals(List.of("ping", "done"), captureResultTypes());
  598 + verify(shellTracker, atLeastOnce()).cancel();
  599 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  600 + assertFalse(Thread.currentThread().isInterrupted());
  601 + }
  602 +
  603 + @Test
  604 + public void testHandleTask_stoppedDuringGenerate() {
  605 + // mock 方法:输出第一段后停止,之后的增量不再写入
  606 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L));
  607 + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any()))
  608 + .thenAnswer(invocation -> {
  609 + Consumer<String> onContent = invocation.getArgument(8);
  610 + onContent.accept("第一段");
  611 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID);
  612 + onContent.accept("第二段");
  613 + return new Ai1LlmChatTool.ChatText("第一段第二段", "");
  614 + });
  615 +
  616 + // 调用
  617 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  618 + // 断言:保留已生成内容,只提交一个终态
  619 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "第一段\n\n> 本条回复已被人工终止".equals(message.getContent())
  620 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  621 + assertEquals(List.of("ping", "message", "done"), captureResultTypes());
  622 + assertFalse(Thread.currentThread().isInterrupted());
  623 + }
  624 +
  625 + @Test
  626 + public void testStop_interruptRunning() {
  627 + // mock 方法:生成期间节点停机
  628 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L));
  629 + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any()))
  630 + .thenAnswer(invocation -> {
  631 + sessionStreamTool.stop();
  632 + return new Ai1LlmChatTool.ChatText("部分", "");
  633 + });
  634 +
  635 + // 调用
  636 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  637 + // 断言:保留内容并标记中断,另写 error,确认后不会被重新生成
  638 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "部分\n\n> 生成任务已中断,请重新发送".equals(message.getContent())
  639 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  640 + assertEquals(List.of("ping", "error", "done"), captureResultTypes());
  641 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  642 + }
  643 +
  644 + @Test
  645 + public void testHandleTask_stoppedBeforeRegister() {
  646 + // mock 方法:查询消息后、登记上下文前节点停机
  647 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenAnswer(invocation -> {
  648 + sessionStreamTool.stop();
  649 + return new Ai1SessionMessageDO().setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  650 + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  651 + });
  652 +
  653 + // 调用
  654 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  655 + // 断言:不调用模型、不装配工具,按中断收尾并确认
  656 + verifyNoInteractions(agentService, llmChatTool, mcpToolFactory);
  657 + verify(stopBucket, never()).isExists();
  658 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "\n\n> 生成任务已中断,请重新发送".equals(message.getContent())
  659 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  660 + assertEquals(List.of("error", "done"), captureResultTypes());
  661 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  662 + assertTrue(getRunContexts().isEmpty());
  663 + assertFalse(Thread.currentThread().isInterrupted());
  664 + }
  665 +
  666 + @Test
  667 + public void testConsumeLoop_stoppedAfterRead() {
  668 + // mock 方法:阻塞读取返回任务时,节点已停机
  669 + when(taskStream.readGroup(anyString(), anyString(), any(StreamReadGroupArgs.class))).thenAnswer(invocation -> {
  670 + sessionStreamTool.stop();
  671 + return Map.of(RECORD_ID, buildTaskFields());
  672 + });
  673 +
  674 + // 调用
  675 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "consumeLoop", "worker-test");
  676 + // 断言:不调用模型,按中断收尾后退出循环
  677 + verifyNoInteractions(agentService, llmChatTool);
  678 + assertEquals(List.of("error", "done"), captureResultTypes());
  679 + verify(taskStream).ack("ai1-session-workers", RECORD_ID);
  680 + assertFalse(Thread.currentThread().isInterrupted());
  681 + }
  682 +
  683 + @Test
  684 + public void testCancel_queued() {
  685 + // mock 方法:会话中有一条生成中的助手消息
  686 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of(MESSAGE_ID));
  687 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(new Ai1SessionMessageDO().setId(MESSAGE_ID)
  688 + .setSessionId(SESSION_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  689 + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus()));
  690 + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenReturn(
  691 + Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成")));
  692 + when(taskLock.isHeldByCurrentThread()).thenReturn(true);
  693 +
  694 + // 调用
  695 + sessionStreamTool.cancel(SESSION_ID);
  696 + // 断言:写入停止意图并广播
  697 + verify(stopBucket).set("1", Duration.ofSeconds(3600));
  698 + verify(stopTopic).publish(String.valueOf(MESSAGE_ID));
  699 + verify(redissonClient, times(1)).getBucket(anyString(), any(Codec.class));
  700 + // 断言:没有 worker 执行,直接收尾
  701 + verify(sessionMessageService).updateSessionMessage(argThat(message -> "已生成\n\n> 本条回复已被人工终止".equals(message.getContent())
  702 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(message.getStatus())));
  703 + verify(sessionService).touchSession(SESSION_ID);
  704 + assertEquals(List.of("done"), captureResultTypes());
  705 + verify(taskLock).unlock();
  706 + verify(sessionMessageService, never()).getSessionMessageListBySessionId(any());
  707 + }
  708 +
  709 + @Test
  710 + public void testCancel_running() {
  711 + // mock 方法:任务正在其他 worker 执行
  712 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of(MESSAGE_ID));
  713 + when(taskLock.tryLock()).thenReturn(false);
  714 +
  715 + // 调用
  716 + sessionStreamTool.cancel(SESSION_ID);
  717 + // 断言:只写意图与通知,由执行节点收尾
  718 + verify(stopBucket).set("1", Duration.ofSeconds(3600));
  719 + verify(stopTopic).publish(String.valueOf(MESSAGE_ID));
  720 + verify(sessionMessageService, never()).updateSessionMessage(any());
  721 + verify(resultStream, never()).add(any());
  722 + }
  723 +
  724 + @Test
  725 + public void testCancel_noGenerating() {
  726 + // mock 方法
  727 + when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenReturn(List.of());
  728 +
  729 + // 调用
  730 + sessionStreamTool.cancel(SESSION_ID);
  731 + // 断言:没有生成中的消息时无副作用
  732 + verifyNoInteractions(stopBucket, stopTopic, taskLock);
  733 + }
  734 +
  735 + @Test
  736 + @Timeout(30)
  737 + public void testCheckStopRequests() throws Exception {
  738 + // 准备参数
  739 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(new Thread(() -> {}), shellTracker, getReapers());
  740 + Map<Long, Ai1SessionStreamTool.RunContext> runContexts = getRunContexts();
  741 + runContexts.put(MESSAGE_ID, context);
  742 + // mock 订阅断开期间写入的停止意图;进程回收一直阻塞
  743 + when(stopBucket.isExists()).thenReturn(true);
  744 + CountDownLatch killing = new CountDownLatch(1);
  745 + CountDownLatch releaseKill = new CountDownLatch(1);
  746 + doAnswer(invocation -> {
  747 + killing.countDown();
  748 + releaseKill.await(10, TimeUnit.SECONDS);
  749 + return null;
  750 + }).when(shellTracker).cancel();
  751 +
  752 + try {
  753 + // 调用:回收阻塞时检查线程仍立即返回
  754 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "checkStopRequests");
  755 + // 断言:回收在独立线程执行
  756 + assertTrue(killing.await(5, TimeUnit.SECONDS));
  757 + assertEquals(Ai1SessionStreamTool.TERMINATED_NOTICE, context.finish());
  758 + } finally {
  759 + releaseKill.countDown();
  760 + awaitReapers();
  761 + }
  762 + }
  763 +
  764 + @Test
  765 + @Timeout(30)
  766 + public void testStop_slowReapNotBlockOthers() throws Exception {
  767 + // 准备参数:两个执行中的生成,第一个的进程回收一直阻塞
  768 + Ai1SkillShellTracker slowTracker = mock(Ai1SkillShellTracker.class);
  769 + Ai1SkillShellTracker fastTracker = mock(Ai1SkillShellTracker.class);
  770 + CountDownLatch killing = new CountDownLatch(1);
  771 + CountDownLatch releaseKill = new CountDownLatch(1);
  772 + CountDownLatch fastReaped = new CountDownLatch(1);
  773 + doAnswer(invocation -> {
  774 + killing.countDown();
  775 + releaseKill.await(10, TimeUnit.SECONDS);
  776 + return null;
  777 + }).when(slowTracker).cancel();
  778 + doAnswer(invocation -> {
  779 + fastReaped.countDown();
  780 + return null;
  781 + }).when(fastTracker).cancel();
  782 + CountDownLatch releaseWorker = new CountDownLatch(1);
  783 + CountDownLatch interrupted = new CountDownLatch(2);
  784 + Thread worker1 = startWorker(releaseWorker, interrupted);
  785 + Thread worker2 = startWorker(releaseWorker, interrupted);
  786 + getRunContexts().put(1L, new Ai1SessionStreamTool.RunContext(worker1, slowTracker, getReapers()));
  787 + getRunContexts().put(2L, new Ai1SessionStreamTool.RunContext(worker2, fastTracker, getReapers()));
  788 + Thread stopping = new Thread(sessionStreamTool::stop);
  789 +
  790 + try {
  791 + // 调用
  792 + stopping.start();
  793 + // 断言:慢回收进行中,两个生成都被中断,另一个回收不受影响;stop 等待回收结束
  794 + assertTrue(killing.await(5, TimeUnit.SECONDS));
  795 + assertTrue(interrupted.await(5, TimeUnit.SECONDS));
  796 + assertTrue(fastReaped.await(5, TimeUnit.SECONDS));
  797 + assertTrue(stopping.isAlive());
  798 + } finally {
  799 + releaseKill.countDown();
  800 + releaseWorker.countDown();
  801 + stopping.join(10_000);
  802 + worker1.join(5000);
  803 + worker2.join(5000);
  804 + }
  805 + assertFalse(stopping.isAlive());
  806 + }
  807 +
  808 + @Test
  809 + @Timeout(30)
  810 + @EnabledOnOs({OS.MAC, OS.LINUX})
  811 + public void testRunContext_cancelReapsWhileWorkerIgnoresInterrupt(@TempDir Path workDir) throws Exception {
  812 + // 准备参数:worker 吞掉中断,模拟模型或工具库不响应中断;Shell 写入自身 PID 后 exec 为 sleep
  813 + CountDownLatch releaseWorker = new CountDownLatch(1);
  814 + Thread worker = new Thread(() -> {
  815 + while (releaseWorker.getCount() > 0) {
  816 + try {
  817 + releaseWorker.await();
  818 + } catch (InterruptedException ignored) {
  819 + // 忽略中断,继续等待
  820 + }
  821 + }
  822 + });
  823 + worker.start();
  824 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(workDir).build());
  825 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(worker, tracker, getReapers());
  826 + ExecHandle handle = tracker.start(ExecSpec.of("echo $$ > pid.tmp && mv pid.tmp pid && exec sleep 60"));
  827 + ProcessHandle process = null;
  828 + try {
  829 + process = awaitDescendant(workDir.resolve("pid"));
  830 + assertTrue(process.isAlive());
  831 +
  832 + // 调用
  833 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE);
  834 + // 断言:worker 仍未返回,在册的后台进程已被回收
  835 + process.onExit().get(10, TimeUnit.SECONDS);
  836 + assertFalse(process.isAlive());
  837 + assertFalse(handle.isAlive());
  838 + assertTrue(worker.isAlive());
  839 + } finally {
  840 + // 只清理本用例启动的线程、进程
  841 + releaseWorker.countDown();
  842 + worker.join(5000);
  843 + awaitReapers();
  844 + if (process != null && process.isAlive()) {
  845 + process.destroyForcibly();
  846 + }
  847 + if (handle.isAlive()) {
  848 + handle.kill();
  849 + }
  850 + }
  851 + }
  852 +
  853 + @Test
  854 + public void testRunContext_cancelThenFinish() throws Exception {
  855 + // 准备参数
  856 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(Thread.currentThread(), shellTracker, getReapers());
  857 +
  858 + // 调用:重复停止只记录第一次
  859 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE);
  860 + context.cancel(Ai1SessionStreamTool.INTERRUPTED_MESSAGE);
  861 + // 断言
  862 + assertTrue(Thread.currentThread().isInterrupted());
  863 + assertThrows(CancellationException.class, context::checkActive);
  864 + assertEquals(Ai1SessionStreamTool.TERMINATED_NOTICE, context.finish());
  865 + assertFalse(Thread.currentThread().isInterrupted());
  866 + // 断言:只启动一次回收
  867 + awaitReapers();
  868 + verify(shellTracker, times(1)).cancel();
  869 + }
  870 +
  871 + @Test
  872 + public void testRunContext_lateCancel() {
  873 + // 准备参数
  874 + Ai1SessionStreamTool.RunContext context = new Ai1SessionStreamTool.RunContext(Thread.currentThread(), shellTracker, getReapers());
  875 + assertNull(context.finish());
  876 +
  877 + // 调用:结束后迟到的停止
  878 + context.cancel(Ai1SessionStreamTool.TERMINATED_NOTICE);
  879 + // 断言:不中断 worker 处理下一个任务,也不启动回收
  880 + assertFalse(Thread.currentThread().isInterrupted());
  881 + assertTrue(getReapers().isEmpty());
  882 + verifyNoInteractions(shellTracker);
  883 + }
  884 +
  885 + @Test
  886 + @Timeout(30)
  887 + public void testStop_waitReaperAfterWorkerExit() throws Exception {
  888 + // mock 方法:生成中启动后台进程,人工停止后进程回收一直阻塞
  889 + ExecBackend backend = mock(ExecBackend.class);
  890 + ExecHandle handle = mock(ExecHandle.class);
  891 + ExecSpec spec = ExecSpec.of("sleep 60");
  892 + when(backend.start(spec)).thenReturn(handle);
  893 + CountDownLatch killing = new CountDownLatch(1);
  894 + CountDownLatch releaseKill = new CountDownLatch(1);
  895 + CountDownLatch killed = new CountDownLatch(1);
  896 + doAnswer(invocation -> {
  897 + killing.countDown();
  898 + releaseKill.await(10, TimeUnit.SECONDS);
  899 + killed.countDown();
  900 + return null;
  901 + }).when(handle).kill();
  902 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(backend);
  903 + when(skillToolFactory.newShellTracker(AGENT_ID)).thenReturn(tracker);
  904 + mockGenerate(new Ai1AgentDO().setId(AGENT_ID).setName("客服").setProviderId(1L).setModelId(2L));
  905 + when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any()))
  906 + .thenAnswer(invocation -> {
  907 + tracker.start(spec);
  908 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "cancelRunContext", MESSAGE_ID);
  909 + // 回收线程认领进程后再返回;worker 线程已被中断,用自旋等待
  910 + long deadline = System.currentTimeMillis() + 5000;
  911 + while (killing.getCount() > 0 && System.currentTimeMillis() < deadline) {
  912 + Thread.onSpinWait();
  913 + }
  914 + return new Ai1LlmChatTool.ChatText("部分", "");
  915 + });
  916 + Thread stopping = new Thread(sessionStreamTool::stop);
  917 +
  918 + try {
  919 + // 调用:worker 结束并移除上下文,回收仍阻塞
  920 + sessionStreamTool.handleTask(RECORD_ID, buildTaskFields());
  921 + assertTrue(killing.await(5, TimeUnit.SECONDS));
  922 + assertTrue(getRunContexts().isEmpty());
  923 + assertEquals(1, killed.getCount());
  924 + // 调用:节点停止
  925 + stopping.start();
  926 + // 断言:stop 仍等待已结束上下文的回收
  927 + stopping.join(1000);
  928 + assertTrue(stopping.isAlive());
  929 + } finally {
  930 + releaseKill.countDown();
  931 + stopping.join(10_000);
  932 + awaitReapers();
  933 + }
  934 + // 断言:回收完成后 stop 返回,worker 未重复回收
  935 + assertFalse(stopping.isAlive());
  936 + assertEquals(0, killed.getCount());
  937 + verify(handle, times(1)).kill();
  938 + assertTrue(getReapers().isEmpty());
  939 + }
  940 +
  941 + @Test
  942 + public void testConsumeLoop_noGroup() {
  943 + // mock 方法:任务队列被删除后消费组丢失,重建后退出循环
  944 + ReflectionTestUtils.setField(sessionStreamTool, "running", true);
  945 + when(taskStream.readGroup(anyString(), anyString(), any(StreamReadGroupArgs.class)))
  946 + .thenThrow(new RedisException("NOGROUP No such key 'ai1:session:tasks' or consumer group"));
  947 + doAnswer(invocation -> {
  948 + ReflectionTestUtils.setField(sessionStreamTool, "running", false);
  949 + return null;
  950 + }).when(taskStream).createGroup(any(StreamCreateGroupArgs.class));
  951 +
  952 + // 调用
  953 + ReflectionTestUtils.invokeMethod(sessionStreamTool, "consumeLoop", "worker-test");
  954 + // 断言
  955 + verify(taskStream).createGroup(any(StreamCreateGroupArgs.class));
  956 + }
  957 +
  958 + @Test
  959 + public void testFailStaleMessage_notStale() {
  960 + // 调用,并断言:未超时时不加锁,避免 worker 误判为在途
  961 + assertFalse(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now())));
  962 + verifyNoInteractions(redissonClient, sessionMessageService);
  963 + }
  964 +
  965 + @Test
  966 + public void testFailStaleMessage_running() {
  967 + // mock 方法:仍有 worker 持有执行锁
  968 + when(taskLock.tryLock()).thenReturn(false);
  969 +
  970 + // 调用,并断言
  971 + assertFalse(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10))));
  972 + verify(sessionMessageService, never()).getSessionMessage(any());
  973 + verify(sessionMessageService, never()).updateSessionMessage(any());
  974 + verify(resultStream, never()).add(any());
  975 + }
  976 +
  977 + @Test
  978 + public void testFailStaleMessage_success() {
  979 + // mock 方法:任务记录已裁剪,没有 worker 执行
  980 + Ai1SessionMessageDO message = buildGeneratingMessage(LocalDateTime.now().minusMinutes(10));
  981 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(message);
  982 + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenReturn(
  983 + Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成")));
  984 +
  985 + // 调用,并断言:与接管任务一致,保留部分结果并标记中断
  986 + assertTrue(sessionStreamTool.failStaleMessage(message));
  987 + verify(sessionMessageService).updateSessionMessage(argThat(updateObj -> "已生成\n\n> 生成任务已中断,请重新发送".equals(updateObj.getContent())
  988 + && Ai1SessionMessageStatusEnum.FAILED.getStatus().equals(updateObj.getStatus())));
  989 + verify(sessionService).touchSession(SESSION_ID);
  990 + assertEquals(List.of("error", "done"), captureResultTypes());
  991 + verify(taskLock).unlock();
  992 + verifyNoInteractions(llmChatTool);
  993 + }
  994 +
  995 + @Test
  996 + public void testFailStaleMessage_completedConcurrently() {
  997 + // mock 方法:加锁前 worker 已完成
  998 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10))
  999 + .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus()));
  1000 +
  1001 + // 调用,并断言:不再改写,归档可继续
  1002 + assertTrue(sessionStreamTool.failStaleMessage(buildGeneratingMessage(LocalDateTime.now().minusMinutes(10))));
  1003 + verify(sessionMessageService, never()).updateSessionMessage(any());
  1004 + verify(resultStream, never()).add(any());
  1005 + verify(taskLock).unlock();
  1006 + }
  1007 +
  1008 + @Test
  1009 + public void testFailStaleMessage_lockLostDuringPartialRead() {
  1010 + // mock 方法:读取部分结果期间执行锁失效,已被新持有者接管
  1011 + Ai1SessionMessageDO message = buildGeneratingMessage(LocalDateTime.now().minusMinutes(10));
  1012 + when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenReturn(message);
  1013 + AtomicBoolean owned = new AtomicBoolean(true);
  1014 + when(taskLock.isHeldByThread(anyLong())).thenAnswer(invocation -> owned.get());
  1015 + when(resultStream.range(StreamRangeArgs.startId(StreamMessageId.MIN).endId(StreamMessageId.MAX))).thenAnswer(invocation -> {
  1016 + owned.set(false);
  1017 + return Map.of(RECORD_ID, Map.of("type", "message", "data", "已生成"));
  1018 + });
  1019 +
  1020 + // 调用,并断言:不改动消息、结果流,也不释放新持有者的锁
  1021 + assertThrows(IllegalStateException.class, () -> sessionStreamTool.failStaleMessage(message));
  1022 + verify(sessionMessageService, never()).updateSessionMessage(any());
  1023 + verify(sessionService, never()).touchSession(any());
  1024 + verify(resultStream, never()).add(any());
  1025 + verify(taskLock, never()).unlock();
  1026 + }
  1027 +
  1028 + // ========== mock 方法 ==========
  1029 +
  1030 + private void mockGenerate(Ai1AgentDO agent) {
  1031 + when(agentService.validateAgentExists(AGENT_ID)).thenReturn(agent);
  1032 + when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType()));
  1033 + when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt()))
  1034 + .thenReturn(new ArrayList<>());
  1035 + }
  1036 +
  1037 + private static Thread startWorker(CountDownLatch release, CountDownLatch interrupted) {
  1038 + Thread worker = new Thread(() -> {
  1039 + try {
  1040 + release.await();
  1041 + } catch (InterruptedException e) {
  1042 + interrupted.countDown();
  1043 + }
  1044 + });
  1045 + worker.start();
  1046 + return worker;
  1047 + }
  1048 +
440 1049 // ========== 测试数据 ==========
441 1050  
442 1051 private static Map<String, String> buildTaskFields() {
... ... @@ -449,6 +1058,33 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
449 1058 return fields;
450 1059 }
451 1060  
  1061 + private static Ai1SessionMessageDO buildGeneratingMessage(LocalDateTime createTime) {
  1062 + Ai1SessionMessageDO message = new Ai1SessionMessageDO().setId(MESSAGE_ID).setSessionId(SESSION_ID)
  1063 + .setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
  1064 + .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  1065 + message.setCreateTime(createTime);
  1066 + return message;
  1067 + }
  1068 +
  1069 + /**
  1070 + * 等待 PID 文件写入,且该进程属于当前 JVM 的子孙进程
  1071 + */
  1072 + private static ProcessHandle awaitDescendant(Path pidFile) throws Exception {
  1073 + long deadline = System.currentTimeMillis() + 5_000;
  1074 + while (System.currentTimeMillis() < deadline) {
  1075 + if (Files.exists(pidFile)) {
  1076 + long pid = Long.parseLong(Files.readString(pidFile).trim());
  1077 + Optional<ProcessHandle> process = ProcessHandle.current().descendants()
  1078 + .filter(descendant -> descendant.pid() == pid).findFirst();
  1079 + if (process.isPresent()) {
  1080 + return process.get();
  1081 + }
  1082 + }
  1083 + Thread.sleep(50);
  1084 + }
  1085 + return fail("未找到本用例启动的 Shell 进程");
  1086 + }
  1087 +
452 1088 // ========== 断言工具 ==========
453 1089  
454 1090 @SuppressWarnings({"unchecked", "rawtypes"})
... ... @@ -463,6 +1099,36 @@ public class Ai1SessionStreamToolTest extends BaseMockitoUnitTest {
463 1099 return types;
464 1100 }
465 1101  
  1102 + @SuppressWarnings({"unchecked", "rawtypes"})
  1103 + private List<String> captureResultData() {
  1104 + ArgumentCaptor<StreamAddArgs> captor = ArgumentCaptor.forClass(StreamAddArgs.class);
  1105 + verify(resultStream, atLeastOnce()).add(captor.capture());
  1106 + List<String> data = new ArrayList<>();
  1107 + for (StreamAddArgs args : captor.getAllValues()) {
  1108 + data.add(String.valueOf(readEntries(args).get("data")));
  1109 + }
  1110 + return data;
  1111 + }
  1112 +
  1113 + @SuppressWarnings("unchecked")
  1114 + private Set<Thread> getReapers() {
  1115 + return (Set<Thread>) ReflectionTestUtils.getField(sessionStreamTool, "reapers");
  1116 + }
  1117 +
  1118 + /**
  1119 + * 等待本用例启动的回收线程结束
  1120 + */
  1121 + private void awaitReapers() throws InterruptedException {
  1122 + for (Thread reaper : getReapers()) {
  1123 + reaper.join(10_000);
  1124 + }
  1125 + }
  1126 +
  1127 + @SuppressWarnings("unchecked")
  1128 + private Map<Long, Ai1SessionStreamTool.RunContext> getRunContexts() {
  1129 + return (Map<Long, Ai1SessionStreamTool.RunContext>) ReflectionTestUtils.getField(sessionStreamTool, "runContexts");
  1130 + }
  1131 +
466 1132 @SuppressWarnings("unchecked")
467 1133 private static Map<String, String> readEntries(StreamAddArgs<String, String> args) {
468 1134 try {
... ...
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;
9 9 import cn.iocoder.yudao.module.ai1.framework.ai.config.YudaoAi1Properties;
10 10 import cn.iocoder.yudao.module.ai1.harness.llm.Ai1LlmChatTool;
11 11 import cn.iocoder.yudao.module.ai1.harness.mcp.Ai1McpToolFactory;
  12 +import cn.iocoder.yudao.module.ai1.harness.memory.Ai1MemoryTool;
12 13 import cn.iocoder.yudao.module.ai1.harness.rag.Ai1RagTool;
  14 +import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillShellTracker;
13 15 import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory;
14 16 import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
15 17 import cn.iocoder.yudao.module.ai1.service.knowledge.Ai1KnowledgeBaseService;
... ... @@ -26,6 +28,7 @@ import org.mockito.Mock;
26 28 import org.redisson.Redisson;
27 29 import org.redisson.api.RScript;
28 30 import org.redisson.api.RStream;
  31 +import org.redisson.api.RTopic;
29 32 import org.redisson.api.RedissonClient;
30 33 import org.redisson.api.stream.StreamAddArgs;
31 34 import org.redisson.api.stream.StreamCreateGroupArgs;
... ... @@ -87,6 +90,10 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
87 90 private Ai1SkillToolFactory skillToolFactory;
88 91 @Mock
89 92 private Ai1RagTool ragTool;
  93 + @Mock
  94 + private Ai1MemoryTool memoryTool;
  95 + @Mock
  96 + private Ai1SkillShellTracker shellTracker;
90 97  
91 98 private RedissonClient firstClient;
92 99 private RedissonClient secondClient;
... ... @@ -114,18 +121,21 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
114 121 message.set(new Ai1SessionMessageDO().setId(MESSAGE_ID).setRole(Ai1SessionMessageRoleEnum.ASSISTANT.getRole())
115 122 .setStatus(Ai1SessionMessageStatusEnum.GENERATING.getStatus()));
116 123  
117   - // mock 业务方法
  124 + // mock 业务方法(排队中停止的用例不调用模型)
118 125 when(sessionMessageService.getSessionMessage(MESSAGE_ID)).thenAnswer(invocation -> message.get());
  126 + lenient().when(sessionMessageService.getGeneratingAssistantMessageIdList(SESSION_ID)).thenAnswer(invocation ->
  127 + Ai1SessionMessageStatusEnum.isGenerating(message.get().getStatus()) ? List.of(MESSAGE_ID) : List.of());
119 128 doAnswer(invocation -> {
120 129 message.set(invocation.getArgument(0));
121 130 return null;
122 131 }).when(sessionMessageService).updateSessionMessage(any());
123   - when(agentService.validateAgentExists(1L)).thenReturn(new Ai1AgentDO().setId(1L).setName("客服")
  132 + lenient().when(skillToolFactory.newShellTracker(1L)).thenReturn(shellTracker);
  133 + lenient().when(agentService.validateAgentExists(1L)).thenReturn(new Ai1AgentDO().setId(1L).setName("客服")
124 134 .setProviderId(1L).setModelId(2L));
125   - when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType()));
126   - when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt()))
  135 + lenient().when(modelService.getModelRespBO(1L, 2L)).thenReturn(new Ai1ModelRespBO().setModelType(Ai1ModelTypeEnum.CHAT.getType()));
  136 + lenient().when(sessionMessageService.getSessionMessageListBySessionIdAndIdLessThan(eq(SESSION_ID), eq(MESSAGE_ID), anyInt()))
127 137 .thenReturn(new ArrayList<>());
128   - when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any()))
  138 + lenient().when(llmChatTool.chat(any(), anyString(), anyList(), anyString(), anyList(), anyList(), anyLong(), any(), any()))
129 139 .thenAnswer(invocation -> {
130 140 Consumer<String> onContent = invocation.getArgument(8);
131 141 onContent.accept("第一段");
... ... @@ -143,6 +153,7 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
143 153 assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
144 154 taskStream.delete();
145 155 resultStream.delete();
  156 + secondClient.getBucket(String.format(SESSION_STOP, MESSAGE_ID)).delete();
146 157 secondClient.getLock(String.format(SESSION_TASK_LOCK, MESSAGE_ID)).forceUnlock();
147 158 firstClient.shutdown();
148 159 secondClient.shutdown();
... ... @@ -195,7 +206,7 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
195 206 secondWorker.handleTask(recordId, reclaimed.get(recordId), true);
196 207  
197 208 // 断言接管只保留部分内容并标记失败,不再次调用模型
198   - assertEquals("第一段", message.get().getContent());
  209 + assertEquals("第一段\n\n> 生成任务已中断,请重新发送", message.get().getContent());
199 210 assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus());
200 211 assertEquals(List.of("ping", "message", "error", "done"), resultTypes());
201 212 assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal());
... ... @@ -204,10 +215,108 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
204 215 // 调用旧节点恢复,断言不追加旧增量或覆盖回填结果
205 216 release.countDown();
206 217 generation.get(5, TimeUnit.SECONDS);
207   - assertEquals("第一段", message.get().getContent());
  218 + assertEquals("第一段\n\n> 生成任务已中断,请重新发送", message.get().getContent());
208 219 assertEquals(List.of("ping", "message", "error", "done"), resultTypes());
209 220 }
210 221  
  222 + @Test
  223 + public void testStop_crossNode() throws Exception {
  224 + // 准备任务,交给第一个节点生成;第一个节点订阅停止通知
  225 + startWithoutWorker(firstWorker);
  226 + StreamMessageId recordId = addTask();
  227 + Map<String, String> fields = taskStream.readGroup(SESSION_TASK_GROUP, "first",
  228 + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId);
  229 + Future<?> generation = executor.submit(() -> firstWorker.handleTask(recordId, fields));
  230 + assertTrue(generated.await(5, TimeUnit.SECONDS));
  231 +
  232 + try {
  233 + // 调用:请求落在第二个节点
  234 + secondWorker.cancel(SESSION_ID);
  235 + generation.get(5, TimeUnit.SECONDS);
  236 +
  237 + // 断言:保留已生成内容,人工终止只写 done,并确认任务
  238 + assertEquals("第一段\n\n> 本条回复已被人工终止", message.get().getContent());
  239 + assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus());
  240 + assertEquals(List.of("ping", "message", "done"), resultTypes());
  241 + assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal());
  242 + verify(shellTracker, atLeastOnce()).cancel();
  243 + // 断言:worker 线程的中断标志已清除,不影响下一轮
  244 + assertFalse(executor.submit(() -> Thread.currentThread().isInterrupted()).get(5, TimeUnit.SECONDS));
  245 + } finally {
  246 + firstWorker.stop();
  247 + }
  248 + }
  249 +
  250 + @Test
  251 + public void testStop_intentWithoutSubscription() throws Exception {
  252 + // 准备任务,交给第一个节点生成;模拟停止通知的订阅已断开
  253 + startWithoutWorker(firstWorker);
  254 + RTopic stopTopic = firstClient.getTopic(SESSION_STOP_TOPIC, StringCodec.INSTANCE);
  255 + stopTopic.removeListener((Integer) ReflectionTestUtils.getField(firstWorker, "stopListenerId"));
  256 + StreamMessageId recordId = addTask();
  257 + Map<String, String> fields = taskStream.readGroup(SESSION_TASK_GROUP, "first",
  258 + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId);
  259 + Future<?> generation = executor.submit(() -> firstWorker.handleTask(recordId, fields));
  260 + assertTrue(generated.await(5, TimeUnit.SECONDS));
  261 +
  262 + try {
  263 + // 调用:停止通知丢失,由定时检查停止意图兜底
  264 + secondWorker.cancel(SESSION_ID);
  265 + generation.get(5, TimeUnit.SECONDS);
  266 +
  267 + // 断言
  268 + assertEquals("第一段\n\n> 本条回复已被人工终止", message.get().getContent());
  269 + assertEquals(List.of("ping", "message", "done"), resultTypes());
  270 + } finally {
  271 + firstWorker.stop();
  272 + }
  273 + }
  274 +
  275 + @Test
  276 + public void testStop_queued() {
  277 + // 准备任务,尚未被任何 worker 读取
  278 + StreamMessageId recordId = addTask();
  279 +
  280 + // 调用:无人执行时由请求节点直接收尾
  281 + secondWorker.cancel(SESSION_ID);
  282 + // 断言
  283 + assertEquals("\n\n> 本条回复已被人工终止", message.get().getContent());
  284 + assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), message.get().getStatus());
  285 + assertEquals(List.of("done"), resultTypes());
  286 +
  287 + // 调用:worker 之后读到该任务
  288 + Map<String, String> fields = taskStream.readGroup(SESSION_TASK_GROUP, "first",
  289 + StreamReadGroupArgs.neverDelivered().count(1)).get(recordId);
  290 + firstWorker.handleTask(recordId, fields);
  291 + // 断言:不调用模型,补齐终态后确认
  292 + verifyNoInteractions(llmChatTool);
  293 + assertEquals(List.of("done"), resultTypes());
  294 + assertEquals(0, taskStream.getPendingInfo(SESSION_TASK_GROUP).getTotal());
  295 + }
  296 +
  297 + @Test
  298 + public void testConsume_missingGroupRecovered() throws Exception {
  299 + // 准备:启动一个 worker 后删除任务队列,消费组随之丢失
  300 + YudaoAi1Properties properties = (YudaoAi1Properties) ReflectionTestUtils.getField(secondWorker, "ai1Properties");
  301 + properties.getSession().getWorker().setCount(1);
  302 + secondWorker.start();
  303 + try {
  304 + taskStream.delete();
  305 + addTask();
  306 + release.countDown();
  307 +
  308 + // 断言:worker 重建消费组后继续消费
  309 + long deadline = System.currentTimeMillis() + 20_000;
  310 + while (Ai1SessionMessageStatusEnum.isGenerating(message.get().getStatus()) && System.currentTimeMillis() < deadline) {
  311 + Thread.sleep(200);
  312 + }
  313 + assertEquals(Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message.get().getStatus());
  314 + assertEquals("第一段第二段", message.get().getContent());
  315 + } finally {
  316 + secondWorker.stop();
  317 + }
  318 + }
  319 +
211 320 // ========== 测试数据 ==========
212 321  
213 322 private RedissonClient createClient() {
... ... @@ -241,9 +350,21 @@ public class Ai1SessionTaskRecoveryIntegrationTest extends BaseMockitoUnitTest {
241 350 ReflectionTestUtils.setField(worker, "mcpToolFactory", mcpToolFactory);
242 351 ReflectionTestUtils.setField(worker, "skillToolFactory", skillToolFactory);
243 352 ReflectionTestUtils.setField(worker, "ragTool", ragTool);
  353 + ReflectionTestUtils.setField(worker, "memoryTool", memoryTool);
  354 + // 直接调用 handleTask 的用例按运行中处理,与 start 之后一致
  355 + ReflectionTestUtils.setField(worker, "running", true);
244 356 return worker;
245 357 }
246 358  
  359 + /**
  360 + * 启动节点但不拉起消费线程,只订阅停止通知、检查停止意图
  361 + */
  362 + private static void startWithoutWorker(Ai1SessionStreamTool worker) {
  363 + YudaoAi1Properties properties = (YudaoAi1Properties) ReflectionTestUtils.getField(worker, "ai1Properties");
  364 + properties.getSession().getWorker().setCount(0);
  365 + worker.start();
  366 + }
  367 +
247 368 private StreamMessageId addTask() {
248 369 return taskStream.add(StreamAddArgs.entries(Map.of("messageId", String.valueOf(MESSAGE_ID),
249 370 "agentId", "1", "sessionId", String.valueOf(SESSION_ID), "content", "你好")));
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/harness/skill/Ai1SkillShellTrackerTest.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.harness.skill;
  2 +
  3 +import cn.iocoder.yudao.framework.test.core.ut.BaseMockitoUnitTest;
  4 +import org.junit.jupiter.api.Test;
  5 +import org.junit.jupiter.api.Timeout;
  6 +import org.junit.jupiter.api.condition.EnabledOnOs;
  7 +import org.junit.jupiter.api.condition.OS;
  8 +import org.junit.jupiter.api.io.TempDir;
  9 +import org.mockito.Mock;
  10 +import org.mockito.stubbing.Answer;
  11 +import org.springaicommunity.agent.common.exec.ExecBackend;
  12 +import org.springaicommunity.agent.common.exec.ExecHandle;
  13 +import org.springaicommunity.agent.common.exec.ExecResult;
  14 +import org.springaicommunity.agent.common.exec.ExecSpec;
  15 +import org.springaicommunity.agent.exec.LocalExecBackend;
  16 +
  17 +import java.nio.file.Files;
  18 +import java.nio.file.Path;
  19 +import java.util.Optional;
  20 +import java.util.concurrent.TimeUnit;
  21 +import java.util.concurrent.atomic.AtomicBoolean;
  22 +
  23 +import static org.junit.jupiter.api.Assertions.*;
  24 +import static org.mockito.Mockito.*;
  25 +
  26 +/**
  27 + * {@link Ai1SkillShellTracker} 的单元测试
  28 + *
  29 + * @author 芋道源码
  30 + */
  31 +public class Ai1SkillShellTrackerTest extends BaseMockitoUnitTest {
  32 +
  33 + private static final ExecSpec SPEC = ExecSpec.of("sleep 60");
  34 +
  35 + @Mock
  36 + private ExecBackend delegate;
  37 + @Mock
  38 + private ExecHandle handle1;
  39 + @Mock
  40 + private ExecHandle handle2;
  41 +
  42 + @Test
  43 + public void testRun_delegate() {
  44 + // mock 方法
  45 + ExecResult result = ExecResult.completed(0, "ok", "");
  46 + when(delegate.run(SPEC)).thenReturn(result);
  47 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate);
  48 +
  49 + // 调用,并断言:同步命令自带超时,不登记
  50 + assertSame(result, tracker.run(SPEC));
  51 + tracker.cancel();
  52 + verifyNoInteractions(handle1);
  53 + }
  54 +
  55 + @Test
  56 + public void testCancel_killAll() {
  57 + // mock 方法:第一个进程回收失败,不影响其他进程
  58 + when(delegate.start(SPEC)).thenReturn(handle1, handle2);
  59 + doThrow(new IllegalStateException("进程已退出")).when(handle1).kill();
  60 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate);
  61 + tracker.start(SPEC);
  62 + tracker.start(SPEC);
  63 +
  64 + // 调用:重复终止幂等
  65 + tracker.cancel();
  66 + tracker.cancel();
  67 + // 断言
  68 + verify(handle1).kill();
  69 + verify(handle2).kill();
  70 + }
  71 +
  72 + @Test
  73 + @Timeout(30)
  74 + public void testCancel_concurrentKillOnce() throws Exception {
  75 + // mock 方法:第一个进程回收期间,另一个线程(如 worker 结束)并发回收
  76 + when(delegate.start(SPEC)).thenReturn(handle1, handle2);
  77 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate);
  78 + tracker.start(SPEC);
  79 + tracker.start(SPEC);
  80 + AtomicBoolean first = new AtomicBoolean(true);
  81 + Answer<Void> concurrentCancel = invocation -> {
  82 + if (first.getAndSet(false)) {
  83 + Thread other = new Thread(tracker::cancel);
  84 + other.start();
  85 + other.join(5000);
  86 + }
  87 + return null;
  88 + };
  89 + doAnswer(concurrentCancel).when(handle1).kill();
  90 + doAnswer(concurrentCancel).when(handle2).kill();
  91 +
  92 + // 调用
  93 + tracker.cancel();
  94 + // 断言:每个进程只回收一次
  95 + verify(handle1, times(1)).kill();
  96 + verify(handle2, times(1)).kill();
  97 + }
  98 +
  99 + @Test
  100 + public void testStart_afterCancel() {
  101 + // 准备参数
  102 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate);
  103 + tracker.cancel();
  104 +
  105 + // 调用,并断言异常:终止后拒绝再启动
  106 + assertThrows(IllegalStateException.class, () -> tracker.start(SPEC));
  107 + verify(delegate, never()).start(any());
  108 + }
  109 +
  110 + @Test
  111 + public void testStart_cancelledDuringStart() {
  112 + // mock 方法:进程启动期间被终止
  113 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(delegate);
  114 + when(delegate.start(SPEC)).thenAnswer(invocation -> {
  115 + tracker.cancel();
  116 + return handle1;
  117 + });
  118 +
  119 + // 调用
  120 + assertSame(handle1, tracker.start(SPEC));
  121 + // 断言:登记后发现已终止,立即回收
  122 + verify(handle1).kill();
  123 + }
  124 +
  125 + @Test
  126 + @Timeout(30)
  127 + @EnabledOnOs({OS.MAC, OS.LINUX})
  128 + public void testCancel_realProcess(@TempDir Path workDir) throws Exception {
  129 + // 准备参数:Shell 写入自身 PID 后 exec 为 sleep,PID 不变
  130 + Ai1SkillShellTracker tracker = new Ai1SkillShellTracker(LocalExecBackend.builder().workingDirectory(workDir).build());
  131 + ExecHandle handle = tracker.start(ExecSpec.of("echo $$ > pid.tmp && mv pid.tmp pid && exec sleep 60"));
  132 + ProcessHandle process = null;
  133 + try {
  134 + process = awaitDescendant(workDir.resolve("pid"));
  135 + assertTrue(process.isAlive());
  136 +
  137 + // 调用
  138 + tracker.cancel();
  139 + // 断言
  140 + process.onExit().get(10, TimeUnit.SECONDS);
  141 + assertFalse(process.isAlive());
  142 + assertFalse(handle.isAlive());
  143 + } finally {
  144 + // 只清理本用例启动的进程
  145 + if (process != null && process.isAlive()) {
  146 + process.destroyForcibly();
  147 + }
  148 + if (handle.isAlive()) {
  149 + handle.kill();
  150 + }
  151 + }
  152 + }
  153 +
  154 + // ========== 测试数据 ==========
  155 +
  156 + /**
  157 + * 等待 PID 文件写入,且该进程属于当前 JVM 的子孙进程
  158 + */
  159 + private static ProcessHandle awaitDescendant(Path pidFile) throws Exception {
  160 + long deadline = System.currentTimeMillis() + 5_000;
  161 + while (System.currentTimeMillis() < deadline) {
  162 + if (Files.exists(pidFile)) {
  163 + long pid = Long.parseLong(Files.readString(pidFile).trim());
  164 + Optional<ProcessHandle> process = ProcessHandle.current().descendants()
  165 + .filter(descendant -> descendant.pid() == pid).findFirst();
  166 + if (process.isPresent()) {
  167 + return process.get();
  168 + }
  169 + }
  170 + Thread.sleep(50);
  171 + }
  172 + return fail("未找到本用例启动的 Shell 进程");
  173 + }
  174 +
  175 +}
... ...
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;
16 16 import org.mockito.InjectMocks;
17 17 import org.mockito.Mock;
18 18 import org.mockito.Spy;
  19 +import org.springaicommunity.agent.common.exec.ExecResult;
  20 +import org.springaicommunity.agent.common.exec.ExecSpec;
19 21 import org.springframework.ai.tool.ToolCallback;
20 22  
21 23 import java.nio.file.Path;
... ... @@ -64,7 +66,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
64 66 mockSkillFiles("echo hi");
65 67  
66 68 // 调用
67   - List<Object> tools = skillToolFactory.buildTools(buildAgent());
  69 + List<Object> tools = buildTools(buildAgent());
68 70 // 断言
69 71 assertEquals(6, tools.size());
70 72 assertInstanceOf(ToolCallback.class, tools.get(0));
... ... @@ -82,8 +84,8 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
82 84 Ai1AgentDO agent = buildAgent();
83 85  
84 86 // 调用
85   - skillToolFactory.buildTools(agent);
86   - skillToolFactory.buildTools(agent);
  87 + buildTools(agent);
  88 + buildTools(agent);
87 89 // 断言
88 90 verify(skillFileService, times(1)).getSkillFileListBySkillId(SKILL_ID);
89 91  
... ... @@ -94,7 +96,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
94 96 mockSkillFiles("echo v2");
95 97  
96 98 // 调用
97   - skillToolFactory.buildTools(agent);
  99 + buildTools(agent);
98 100 // 断言
99 101 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID);
100 102 assertEquals("echo v2", FileUtil.readUtf8String(agentRoot().resolve("pdf-reader/scripts/run.sh").toFile()));
... ... @@ -107,13 +109,13 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
107 109 mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0));
108 110 mockSkillFiles("echo hi");
109 111 Ai1AgentDO agent = buildAgent();
110   - skillToolFactory.buildTools(agent);
  112 + buildTools(agent);
111 113  
112 114 // mock 物化目录被删除
113 115 FileUtil.del(agentRoot());
114 116  
115 117 // 调用
116   - skillToolFactory.buildTools(agent);
  118 + buildTools(agent);
117 119 // 断言
118 120 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID);
119 121 assertTrue(FileUtil.exist(agentRoot().resolve("pdf-reader/SKILL.md").toFile()));
... ... @@ -125,7 +127,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
125 127 mockSkill(CommonStatusEnum.DISABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0));
126 128  
127 129 // 调用
128   - List<Object> tools = skillToolFactory.buildTools(buildAgent());
  130 + List<Object> tools = buildTools(buildAgent());
129 131 // 断言只保留执行工具,不物化 SKILL 文件
130 132 assertEquals(5, tools.size());
131 133 assertTrue(tools.stream().noneMatch(ToolCallback.class::isInstance));
... ... @@ -136,7 +138,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
136 138 @Test
137 139 public void testBuildTools_noSkill() {
138 140 // 调用,并断言
139   - assertTrue(skillToolFactory.buildTools(new Ai1AgentDO().setId(AGENT_ID).setSkillIds(List.of())).isEmpty());
  141 + assertTrue(buildTools(new Ai1AgentDO().setId(AGENT_ID).setSkillIds(List.of())).isEmpty());
140 142 verify(skillService, never()).getSkillList(anyCollection());
141 143 }
142 144  
... ... @@ -146,7 +148,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
146 148 mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0));
147 149 mockSkillFiles("echo hi");
148 150 Ai1AgentDO agent = buildAgent();
149   - skillToolFactory.buildTools(agent);
  151 + buildTools(agent);
150 152 assertTrue(FileUtil.exist(agentRoot().toFile()));
151 153  
152 154 // 调用
... ... @@ -155,7 +157,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
155 157 assertFalse(FileUtil.exist(agentRoot().toFile()));
156 158  
157 159 // 调用
158   - skillToolFactory.buildTools(agent);
  160 + buildTools(agent);
159 161 // 断言重新物化
160 162 verify(skillFileService, times(2)).getSkillFileListBySkillId(SKILL_ID);
161 163 }
... ... @@ -176,7 +178,7 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
176 178 }
177 179  
178 180 // 调用
179   - ToolCallback tool = (ToolCallback) skillToolFactory.buildTools(agent).get(0);
  181 + ToolCallback tool = (ToolCallback) buildTools(agent).get(0);
180 182 String result = tool.call("{\"command\":\"same\"}");
181 183 // 断言
182 184 assertTrue(result.contains("first-content"));
... ... @@ -185,14 +187,32 @@ public class Ai1SkillToolFactoryTest extends BaseMockitoUnitTest {
185 187  
186 188 // 调用
187 189 skillToolFactory.evict(AGENT_ID);
188   - ToolCallback rebuiltTool = (ToolCallback) skillToolFactory.buildTools(agent).get(0);
  190 + ToolCallback rebuiltTool = (ToolCallback) buildTools(agent).get(0);
189 191 String rebuiltResult = rebuiltTool.call("{\"command\":\"same\"}");
190 192 // 断言
191 193 assertEquals(result, rebuiltResult);
192 194 }
193 195  
  196 + @Test
  197 + public void testNewShellTracker_workingDirectory() throws Exception {
  198 + // mock skillService 和 skillFileService 的方法
  199 + mockSkill(CommonStatusEnum.ENABLE.getStatus(), LocalDateTime.of(2026, 10, 1, 10, 0));
  200 + mockSkillFiles("echo hi");
  201 + buildTools(buildAgent());
  202 +
  203 + // 调用
  204 + ExecResult result = skillToolFactory.newShellTracker(AGENT_ID).run(ExecSpec.of("pwd"));
  205 + // 断言:Shell 命令限定在 Agent 物化根目录执行
  206 + assertEquals(0, result.exitCode());
  207 + assertEquals(agentRoot().toRealPath(), Path.of(result.stdout().trim()).toRealPath());
  208 + }
  209 +
194 210 // ========== 测试数据 ==========
195 211  
  212 + private List<Object> buildTools(Ai1AgentDO agent) {
  213 + return skillToolFactory.buildTools(agent, skillToolFactory.newShellTracker(agent.getId()));
  214 + }
  215 +
196 216 private Path agentRoot() {
197 217 return tempDir.resolve("agent_" + AGENT_ID);
198 218 }
... ...
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;
7 7 import cn.iocoder.yudao.module.ai1.dal.mysql.agent.Ai1AgentMapper;
8 8 import cn.iocoder.yudao.module.ai1.enums.model.Ai1ModelTypeEnum;
9 9 import cn.iocoder.yudao.module.ai1.harness.skill.Ai1SkillToolFactory;
  10 +import cn.iocoder.yudao.module.ai1.service.memory.Ai1MemoryService;
10 11 import cn.iocoder.yudao.module.ai1.service.session.Ai1SessionService;
11 12 import cn.iocoder.yudao.module.ai1.service.model.Ai1ModelService;
12 13 import cn.iocoder.yudao.module.ai1.service.model.bo.Ai1ModelRespBO;
... ... @@ -45,6 +46,8 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest {
45 46 @MockitoBean
46 47 private Ai1SessionService sessionService;
47 48 @MockitoBean
  49 + private Ai1MemoryService memoryService;
  50 + @MockitoBean
48 51 private Ai1SkillToolFactory skillToolFactory;
49 52  
50 53 @Test
... ... @@ -123,6 +126,23 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest {
123 126 }
124 127  
125 128 @Test
  129 + public void testGetAgentListByStatus_all() {
  130 + // mock 数据:开启、关闭、已删除的 Agent 各一个
  131 + mockChatModel();
  132 + Long enabledId = agentService.createAgent(buildAgentSaveReqVO());
  133 + agentService.updateAgentStatus(enabledId, CommonStatusEnum.ENABLE.getStatus());
  134 + Long disabledId = agentService.createAgent(buildAgentSaveReqVO());
  135 + agentMapper.deleteById(agentService.createAgent(buildAgentSaveReqVO()));
  136 +
  137 + // 调用:状态为空时返回全部状态
  138 + List<Ai1AgentDO> list = agentService.getAgentListByStatus(null);
  139 + // 断言:按编号升序,不含已删除的 Agent
  140 + assertEquals(List.of(enabledId, disabledId), list.stream().map(Ai1AgentDO::getId).toList());
  141 + assertEquals(CommonStatusEnum.ENABLE.getStatus(), list.get(0).getStatus());
  142 + assertEquals(CommonStatusEnum.DISABLE.getStatus(), list.get(1).getStatus());
  143 + }
  144 +
  145 + @Test
126 146 public void testDeleteAgent_cascade() {
127 147 // mock 数据
128 148 mockChatModel();
... ... @@ -133,10 +153,29 @@ public class Ai1AgentServiceImplTest extends BaseDbUnitTest {
133 153 // 断言
134 154 assertNull(agentMapper.selectById(id));
135 155 verify(sessionService).deleteSessionListByAgentIds(Collections.singletonList(id));
  156 + verify(memoryService).deleteMemoryListByAgentIds(Collections.singletonList(id));
136 157 verify(skillToolFactory).evict(id);
137 158 }
138 159  
139 160 @Test
  161 + public void testCreateAgent_memoryEnabled() {
  162 + // mock 数据
  163 + mockChatModel();
  164 +
  165 + // 调用
  166 + Long defaultId = agentService.createAgent(buildAgentSaveReqVO());
  167 + Long enabledId = agentService.createAgent(buildAgentSaveReqVO().setMemoryEnabled(true));
  168 + // 断言:默认关闭
  169 + assertFalse(agentMapper.selectById(defaultId).getMemoryEnabled());
  170 + assertTrue(agentMapper.selectById(enabledId).getMemoryEnabled());
  171 +
  172 + // 调用:关闭
  173 + agentService.updateAgent(buildAgentSaveReqVO().setId(enabledId).setMemoryEnabled(false));
  174 + // 断言
  175 + assertFalse(agentMapper.selectById(enabledId).getMemoryEnabled());
  176 + }
  177 +
  178 + @Test
140 179 public void testGetAgentCountByKnowledgeBaseIdsAndSkillIds() {
141 180 // mock 数据
142 181 mockChatModel();
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/memory/Ai1MemoryServiceImplTest.java 0 → 100644
  1 +package cn.iocoder.yudao.module.ai1.service.memory;
  2 +
  3 +import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
  4 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
  5 +import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest;
  6 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryPageReqVO;
  7 +import cn.iocoder.yudao.module.ai1.controller.admin.memory.vo.Ai1MemoryUpdateReqVO;
  8 +import cn.iocoder.yudao.module.ai1.dal.dataobject.agent.Ai1AgentDO;
  9 +import cn.iocoder.yudao.module.ai1.dal.dataobject.memory.Ai1MemoryDO;
  10 +import cn.iocoder.yudao.module.ai1.dal.mysql.memory.Ai1MemoryMapper;
  11 +import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
  12 +import jakarta.annotation.Resource;
  13 +import org.junit.jupiter.api.Test;
  14 +import org.springframework.context.annotation.Import;
  15 +import org.springframework.test.context.bean.override.mockito.MockitoBean;
  16 +
  17 +import java.util.Arrays;
  18 +import java.util.Collections;
  19 +import java.util.List;
  20 +
  21 +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException;
  22 +import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.MEMORY_NOT_EXISTS;
  23 +import static org.junit.jupiter.api.Assertions.*;
  24 +import static org.mockito.Mockito.when;
  25 +
  26 +/**
  27 + * {@link Ai1MemoryServiceImpl} 的单元测试
  28 + *
  29 + * @author 芋道源码
  30 + */
  31 +@Import(Ai1MemoryServiceImpl.class)
  32 +public class Ai1MemoryServiceImplTest extends BaseDbUnitTest {
  33 +
  34 + @Resource
  35 + private Ai1MemoryServiceImpl memoryService;
  36 +
  37 + @Resource
  38 + private Ai1MemoryMapper memoryMapper;
  39 +
  40 + @MockitoBean
  41 + private Ai1AgentService agentService;
  42 +
  43 + @Test
  44 + public void testUpdateMemory_success() {
  45 + // mock 数据
  46 + Ai1MemoryDO memory = insertMemory(1L, 100L);
  47 +
  48 + // 调用
  49 + memoryService.updateMemory(new Ai1MemoryUpdateReqVO().setId(memory.getId()).setContent("## 用户画像\n- 新内容"));
  50 + // 断言:正文更新,水位线不变
  51 + Ai1MemoryDO result = memoryMapper.selectById(memory.getId());
  52 + assertEquals("## 用户画像\n- 新内容", result.getContent());
  53 + assertEquals(10L, result.getLastMessageId());
  54 + }
  55 +
  56 + @Test
  57 + public void testUpdateMemory_notExists() {
  58 + // 调用,并断言异常
  59 + assertServiceException(() -> memoryService.updateMemory(new Ai1MemoryUpdateReqVO().setId(1L).setContent("内容")),
  60 + MEMORY_NOT_EXISTS);
  61 + }
  62 +
  63 + @Test
  64 + public void testDeleteMemoryListByIds_success() {
  65 + // mock 数据
  66 + Ai1MemoryDO memory1 = insertMemory(1L, 100L);
  67 + Ai1MemoryDO memory2 = insertMemory(1L, 200L);
  68 + Ai1MemoryDO memory3 = insertMemory(2L, 100L);
  69 +
  70 + // 调用
  71 + memoryService.deleteMemoryListByIds(Arrays.asList(memory1.getId(), memory2.getId()));
  72 + // 断言:逻辑删除后不再可见,同一用户可以重新生成
  73 + assertNull(memoryMapper.selectById(memory1.getId()));
  74 + assertNull(memoryMapper.selectById(memory2.getId()));
  75 + assertNotNull(memoryMapper.selectById(memory3.getId()));
  76 + assertNotNull(insertMemory(1L, 100L).getId());
  77 + }
  78 +
  79 + @Test
  80 + public void testDeleteMemoryListByIds_notExists() {
  81 + // mock 数据
  82 + Ai1MemoryDO memory = insertMemory(1L, 100L);
  83 +
  84 + // 调用,并断言异常
  85 + assertServiceException(() -> memoryService.deleteMemoryListByIds(Arrays.asList(memory.getId(), memory.getId() + 1)),
  86 + MEMORY_NOT_EXISTS);
  87 + // 断言
  88 + assertNotNull(memoryMapper.selectById(memory.getId()));
  89 + }
  90 +
  91 + @Test
  92 + public void testDeleteMemoryListByAgentIds() {
  93 + // mock 数据
  94 + Ai1MemoryDO memory1 = insertMemory(1L, 100L);
  95 + Ai1MemoryDO memory2 = insertMemory(2L, 100L);
  96 + Ai1MemoryDO memory3 = insertMemory(3L, 100L);
  97 +
  98 + // 调用
  99 + memoryService.deleteMemoryListByAgentIds(Arrays.asList(1L, 2L));
  100 + memoryService.deleteMemoryListByAgentIds(Collections.emptyList());
  101 + // 断言
  102 + assertNull(memoryMapper.selectById(memory1.getId()));
  103 + assertNull(memoryMapper.selectById(memory2.getId()));
  104 + assertNotNull(memoryMapper.selectById(memory3.getId()));
  105 + }
  106 +
  107 + @Test
  108 + public void testGetMemoryPage() {
  109 + // mock 数据
  110 + Ai1MemoryDO memory1 = insertMemory(1L, 100L);
  111 + Ai1MemoryDO memory2 = insertMemory(1L, 200L);
  112 + Ai1MemoryDO memory3 = insertMemory(2L, 100L);
  113 + memoryService.deleteMemory(insertMemory(1L, 300L).getId());
  114 +
  115 + // 调用:不传条件,按编号倒序,不含已删除的记忆
  116 + PageResult<Ai1MemoryDO> pageResult = memoryService.getMemoryPage(new Ai1MemoryPageReqVO());
  117 + // 断言
  118 + assertEquals(List.of(memory3.getId(), memory2.getId(), memory1.getId()),
  119 + pageResult.getList().stream().map(Ai1MemoryDO::getId).toList());
  120 +
  121 + // 调用:按 Agent、用户筛选
  122 + pageResult = memoryService.getMemoryPage(new Ai1MemoryPageReqVO().setAgentId(1L).setUserId(100L));
  123 + // 断言
  124 + assertEquals(List.of(memory1.getId()), pageResult.getList().stream().map(Ai1MemoryDO::getId).toList());
  125 + }
  126 +
  127 + @Test
  128 + public void testGetMemoryByAgentIdAndUserId() {
  129 + // mock 数据
  130 + Ai1MemoryDO memory = insertMemory(1L, 100L);
  131 + insertMemory(1L, 200L);
  132 +
  133 + // 调用,并断言
  134 + assertEquals(memory.getId(), memoryService.getMemoryByAgentIdAndUserId(1L, 100L).getId());
  135 + assertNull(memoryService.getMemoryByAgentIdAndUserId(2L, 100L));
  136 + }
  137 +
  138 + @Test
  139 + public void testGetMemoryByAgentIdAndUserId_duplicate() {
  140 + // mock 数据:极端并发下同一用户产生了重复记录
  141 + insertMemory(1L, 100L);
  142 + insertMemory(1L, 100L);
  143 +
  144 + // 调用,并断言:不报错,返回其中一条
  145 + assertNotNull(memoryService.getMemoryByAgentIdAndUserId(1L, 100L));
  146 + }
  147 +
  148 + @Test
  149 + public void testCreateMemoryIfAgentExists() {
  150 + // mock 方法:Agent 1 存在,Agent 2 已删除
  151 + when(agentService.getAgent(1L)).thenReturn(Ai1AgentDO.builder().id(1L).name("客服")
  152 + .status(CommonStatusEnum.ENABLE.getStatus()).build());
  153 +
  154 + // 调用,并断言:Agent 已删除时不创建
  155 + assertFalse(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(2L)
  156 + .userId(100L).content("迟到的归档结果").lastMessageId(20L).build()));
  157 + assertNull(memoryMapper.selectByAgentIdAndUserId(2L, 100L));
  158 + // 调用,并断言:Agent 存在时创建
  159 + assertTrue(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L)
  160 + .userId(100L).content("用户偏好").lastMessageId(20L).build()));
  161 + assertEquals("用户偏好", memoryMapper.selectByAgentIdAndUserId(1L, 100L).getContent());
  162 + // 调用,并断言:已有未删除的记忆时不重复创建
  163 + assertFalse(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L)
  164 + .userId(100L).content("重复归档").lastMessageId(30L).build()));
  165 + assertEquals(1, memoryMapper.selectCount(Ai1MemoryDO::getAgentId, 1L));
  166 + }
  167 +
  168 + @Test
  169 + public void testCreateMemoryIfAgentExists_afterDelete() {
  170 + // mock 数据:已删除的历史记忆
  171 + when(agentService.getAgent(1L)).thenReturn(Ai1AgentDO.builder().id(1L).name("客服")
  172 + .status(CommonStatusEnum.ENABLE.getStatus()).build());
  173 + Ai1MemoryDO old = insertMemory(1L, 100L);
  174 + memoryService.deleteMemory(old.getId());
  175 +
  176 + // 调用:删除后按历史消息重新生成
  177 + assertTrue(memoryService.createMemoryIfAgentExists(Ai1MemoryDO.builder().agentId(1L)
  178 + .userId(100L).content("重建记忆").lastMessageId(5L).build()));
  179 + // 断言:生成新的记录与水位线,历史记录不影响按用户查询
  180 + Ai1MemoryDO memory = memoryService.getMemoryByAgentIdAndUserId(1L, 100L);
  181 + assertNotEquals(old.getId(), memory.getId());
  182 + assertEquals("重建记忆", memory.getContent());
  183 + assertEquals(5L, memory.getLastMessageId());
  184 + // 断言:旧快照的归档写入不能复活已删除的记录
  185 + assertFalse(memoryService.updateMemoryIfUnchanged(old, "迟到的归档结果", 30L));
  186 + assertNull(memoryService.getMemory(old.getId()));
  187 + }
  188 +
  189 + @Test
  190 + public void testUpdateMemoryIfUnchanged() {
  191 + // mock 数据
  192 + Ai1MemoryDO original = insertMemory(1L, 100L);
  193 +
  194 + // 调用,并断言:首次写入成功,按旧快照再次写入失败
  195 + assertTrue(memoryService.updateMemoryIfUnchanged(original, "新记忆", 20L));
  196 + assertFalse(memoryService.updateMemoryIfUnchanged(original, "过期结果", 30L));
  197 + Ai1MemoryDO memory = memoryMapper.selectById(original.getId());
  198 + assertEquals("新记忆", memory.getContent());
  199 + assertEquals(20L, memory.getLastMessageId());
  200 + }
  201 +
  202 + @Test
  203 + public void testUpdateMemoryIfUnchanged_edited() {
  204 + // mock 数据:归档读取后,人工只在末尾加了空格,水位线不变
  205 + Ai1MemoryDO original = insertMemory(1L, 100L);
  206 + memoryMapper.updateById(new Ai1MemoryDO().setId(original.getId()).setContent("## 用户画像\n- 旧内容 "));
  207 +
  208 + // 调用,并断言:放弃写入,保留人工修改
  209 + assertFalse(memoryService.updateMemoryIfUnchanged(original, "归档结果", 20L));
  210 + Ai1MemoryDO memory = memoryMapper.selectById(original.getId());
  211 + assertEquals("## 用户画像\n- 旧内容 ", memory.getContent());
  212 + assertEquals(10L, memory.getLastMessageId());
  213 + }
  214 +
  215 + // ========== 测试数据 ==========
  216 +
  217 + private Ai1MemoryDO insertMemory(Long agentId, Long userId) {
  218 + Ai1MemoryDO memory = Ai1MemoryDO.builder().agentId(agentId).userId(userId)
  219 + .content("## 用户画像\n- 旧内容").lastMessageId(10L).build();
  220 + memoryMapper.insert(memory);
  221 + return memory;
  222 + }
  223 +
  224 +}
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionMessageServiceImplTest.java
1 1 package cn.iocoder.yudao.module.ai1.service.session;
2 2  
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
3 4 import cn.iocoder.yudao.framework.tenant.core.context.TenantContextHolder;
4 5 import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest;
  6 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessagePageReqVO;
5 7 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.message.Ai1SessionMessageSendReqVO;
6 8 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
7 9 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionMessageDO;
8 10 import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMessageMapper;
  11 +import cn.iocoder.yudao.module.ai1.dal.redis.memory.Ai1MemoryLockRedisDAO;
9 12 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageRoleEnum;
10 13 import cn.iocoder.yudao.module.ai1.enums.session.Ai1SessionMessageStatusEnum;
11 14 import cn.iocoder.yudao.module.ai1.service.agent.Ai1AgentService;
... ... @@ -17,6 +20,9 @@ import org.junit.jupiter.api.BeforeEach;
17 20 import org.junit.jupiter.api.Test;
18 21 import org.junit.jupiter.params.ParameterizedTest;
19 22 import org.junit.jupiter.params.provider.ValueSource;
  23 +import org.mockito.InOrder;
  24 +import org.redisson.api.RLock;
  25 +import org.redisson.api.RedissonClient;
20 26 import org.redisson.client.RedisConnectionException;
21 27 import org.redisson.client.RedisTimeoutException;
22 28 import org.springframework.context.annotation.Import;
... ... @@ -24,9 +30,11 @@ import org.springframework.test.context.bean.override.mockito.MockitoBean;
24 30 import org.springframework.test.context.bean.override.mockito.MockitoSpyBean;
25 31 import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
26 32  
  33 +import java.util.ArrayList;
27 34 import java.util.List;
28 35  
29 36 import static cn.iocoder.yudao.framework.common.exception.util.ServiceExceptionUtil.exception;
  37 +import static cn.iocoder.yudao.framework.test.core.util.AssertUtils.assertServiceException;
30 38 import static cn.iocoder.yudao.module.ai1.enums.Ai1ErrorCodeConstants.SESSION_NOT_EXISTS;
31 39 import static org.junit.jupiter.api.Assertions.*;
32 40 import static org.mockito.ArgumentMatchers.*;
... ... @@ -37,7 +45,7 @@ import static org.mockito.Mockito.*;
37 45 *
38 46 * @author 芋道源码
39 47 */
40   -@Import(Ai1SessionMessageServiceImpl.class)
  48 +@Import({Ai1SessionMessageServiceImpl.class, Ai1MemoryLockRedisDAO.class})
41 49 public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
42 50  
43 51 private static final Long TENANT_ID = 1L;
... ... @@ -55,11 +63,17 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
55 63 private Ai1AgentService agentService;
56 64 @MockitoBean
57 65 private Ai1SessionStreamTool sessionStreamTool;
  66 + @MockitoBean
  67 + private RedissonClient redissonClient;
  68 +
  69 + private final RLock snapshotLock = mock(RLock.class);
58 70  
59 71 @BeforeEach
60 72 public void before() {
61 73 // 设置租户上下文
62 74 TenantContextHolder.setTenantId(TENANT_ID);
  75 + // mock 记忆快照锁
  76 + when(redissonClient.getLock(anyString())).thenReturn(snapshotLock);
63 77 }
64 78  
65 79 @AfterEach
... ... @@ -100,6 +114,71 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
100 114 }
101 115  
102 116 @Test
  117 + public void testSendSessionMessageStream_snapshotLockUntilCommitted() {
  118 + // mock sessionService 的方法
  119 + when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO()
  120 + .setId(10L).setAgentId(1L).setUserId(USER_ID).setTitle(Ai1SessionDO.TITLE_DEFAULT));
  121 + // mock 释放快照锁时,本轮消息已提交
  122 + List<Integer> committedCounts = new ArrayList<>();
  123 + doAnswer(invocation -> {
  124 + committedCounts.add(sessionMessageMapper.selectListBySessionId(10L).size());
  125 + return null;
  126 + }).when(snapshotLock).unlock();
  127 +
  128 + // 调用
  129 + sessionMessageService.sendSessionMessageStream(USER_ID,
  130 + new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好"));
  131 + // 断言:按 Agent、用户加锁,落库后才释放,再投递任务
  132 + verify(redissonClient).getLock("ai1:memory:snapshot-lock:1:" + USER_ID);
  133 + assertEquals(List.of(2), committedCounts);
  134 + InOrder inOrder = inOrder(snapshotLock, sessionStreamTool);
  135 + inOrder.verify(snapshotLock).lock();
  136 + inOrder.verify(snapshotLock).unlock();
  137 + inOrder.verify(sessionStreamTool).submit(anyLong(), eq(1L), eq(10L), eq("你好"));
  138 + }
  139 +
  140 + @Test
  141 + public void testStopSessionMessageMy_success() {
  142 + // 调用
  143 + sessionMessageService.stopSessionMessageMy(USER_ID, 10L);
  144 + // 断言
  145 + InOrder inOrder = inOrder(sessionService, sessionStreamTool);
  146 + inOrder.verify(sessionService).validateSessionMy(USER_ID, 10L);
  147 + inOrder.verify(sessionStreamTool).cancel(10L);
  148 + }
  149 +
  150 + @Test
  151 + public void testStopSessionMessageMy_notOwner() {
  152 + // mock sessionService 的方法
  153 + when(sessionService.validateSessionMy(USER_ID, 10L)).thenThrow(exception(SESSION_NOT_EXISTS));
  154 +
  155 + // 调用,并断言异常
  156 + assertServiceException(() -> sessionMessageService.stopSessionMessageMy(USER_ID, 10L), SESSION_NOT_EXISTS);
  157 + // 断言
  158 + verify(sessionStreamTool, never()).cancel(anyLong());
  159 + }
  160 +
  161 + @Test
  162 + public void testGetSessionMessagePage() {
  163 + // mock 数据
  164 + Ai1SessionMessageDO message1 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole());
  165 + Ai1SessionMessageDO message2 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole());
  166 + Ai1SessionMessageDO message3 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole());
  167 + insertSessionMessage(11L, Ai1SessionMessageRoleEnum.USER.getRole());
  168 + // 准备参数
  169 + Ai1SessionMessagePageReqVO reqVO = new Ai1SessionMessagePageReqVO().setSessionId(10L);
  170 + reqVO.setPageSize(2);
  171 +
  172 + // 调用
  173 + PageResult<Ai1SessionMessageDO> pageResult = sessionMessageService.getSessionMessagePage(reqVO);
  174 + // 断言:按编号倒序
  175 + assertEquals(3, pageResult.getTotal());
  176 + assertEquals(List.of(message3.getId(), message2.getId()),
  177 + pageResult.getList().stream().map(Ai1SessionMessageDO::getId).toList());
  178 + assertNotEquals(message1.getId(), pageResult.getList().get(1).getId());
  179 + }
  180 +
  181 + @Test
103 182 public void testSendSessionMessageStream_submitFailed() {
104 183 // mock sessionService 和 sessionStreamTool 的方法
105 184 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO()
... ... @@ -122,10 +201,11 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
122 201  
123 202 @ParameterizedTest
124 203 @ValueSource(ints = {1, 2})
125   - public void testSendSessionMessageStream_submitTimeoutKeepsWorkerResult(int status) {
  204 + public void testSendSessionMessageStream_submitTimeoutAfterWorkerCompleted(int status) {
126 205 // mock sessionService 和 sessionStreamTool 的方法
127 206 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO()
128 207 .setId(10L).setAgentId(1L).setTitle(Ai1SessionDO.TITLE_DEFAULT));
  208 + // mock 投递实际成功,worker 在投递超时前已回填终态
129 209 doAnswer(invocation -> {
130 210 Long messageId = invocation.getArgument(0);
131 211 sessionMessageMapper.updateById(new Ai1SessionMessageDO().setId(messageId)
... ... @@ -136,7 +216,7 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
136 216 // 调用
137 217 sessionMessageService.sendSessionMessageStream(USER_ID,
138 218 new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好"));
139   - // 断言已有终态和内容不被补偿覆盖
  219 + // 断言:补偿只更新生成中的消息,已有终态和内容不被覆盖
140 220 Ai1SessionMessageDO assistant = sessionMessageMapper.selectListBySessionId(10L).get(1);
141 221 assertEquals(status, assistant.getStatus());
142 222 assertEquals("worker 已回填", assistant.getContent());
... ... @@ -145,28 +225,6 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
145 225 }
146 226  
147 227 @Test
148   - public void testSendSessionMessageStream_workerCanCompleteAfterSubmitTimeout() {
149   - // mock sessionService 和 sessionStreamTool 的方法
150   - when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO()
151   - .setId(10L).setAgentId(1L).setTitle(Ai1SessionDO.TITLE_DEFAULT));
152   - doThrow(new RedisTimeoutException("投递响应丢失")).when(sessionStreamTool)
153   - .submit(anyLong(), eq(1L), eq(10L), eq("你好"));
154   -
155   - // 调用
156   - sessionMessageService.sendSessionMessageStream(USER_ID,
157   - new Ai1SessionMessageSendReqVO().setSessionId(10L).setContent("你好"));
158   - Ai1SessionMessageDO assistant = sessionMessageMapper.selectListBySessionId(10L).get(1);
159   - assertEquals(Ai1SessionMessageStatusEnum.FAILED.getStatus(), assistant.getStatus());
160   - // 调用 worker 回填,模拟 Redis 实际写入后响应丢失
161   - sessionMessageService.updateSessionMessage(new Ai1SessionMessageDO().setId(assistant.getId())
162   - .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus()).setContent("稍后生成成功"));
163   - // 断言后续完成结果仍可落库
164   - Ai1SessionMessageDO result = sessionMessageMapper.selectById(assistant.getId());
165   - assertEquals(Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), result.getStatus());
166   - assertEquals("稍后生成成功", result.getContent());
167   - }
168   -
169   - @Test
170 228 public void testSendSessionMessageStream_openFailedKeepsGenerating() {
171 229 // mock sessionService 和 sessionStreamTool 的方法
172 230 when(sessionService.validateSessionMy(USER_ID, 10L)).thenReturn(new Ai1SessionDO()
... ... @@ -322,11 +380,71 @@ public class Ai1SessionMessageServiceImplTest extends BaseDbUnitTest {
322 380 assertNotEquals(message1.getId(), list.get(1).getId());
323 381 }
324 382  
  383 + @Test
  384 + public void testGetSessionMessageListBySessionIdsAndStatusAndIdBetween_safePrefix() {
  385 + // mock 数据:同一用户两个会话,第一个会话的助手消息仍在生成
  386 + String user = Ai1SessionMessageRoleEnum.USER.getRole();
  387 + String assistant = Ai1SessionMessageRoleEnum.ASSISTANT.getRole();
  388 + Ai1SessionMessageDO message1 = insertSessionMessage(10L, user);
  389 + Ai1SessionMessageDO message2 = insertSessionMessage(10L, assistant, Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  390 + Ai1SessionMessageDO message3 = insertSessionMessage(11L, user);
  391 + insertSessionMessage(12L, user);
  392 + insertSessionMessage(11L, assistant, Ai1SessionMessageStatusEnum.FAILED.getStatus());
  393 + List<Long> sessionIds = List.of(10L, 11L);
  394 +
  395 + // 调用:只读取第一条生成中消息之前的已完成消息
  396 + Ai1SessionMessageDO generating = sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(sessionIds,
  397 + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), 0L);
  398 + List<Ai1SessionMessageDO> first = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds,
  399 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), 0L, generating.getId(), 100);
  400 + // 断言
  401 + assertEquals(message2.getId(), generating.getId());
  402 + assertEquals(List.of(message1.getId()), first.stream().map(Ai1SessionMessageDO::getId).toList());
  403 +
  404 + // mock 生成结束
  405 + sessionMessageMapper.updateById(new Ai1SessionMessageDO().setId(message2.getId())
  406 + .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus()));
  407 + // 调用:从水位线之后读取,并按条数截断
  408 + assertNull(sessionMessageService.getFirstSessionMessageBySessionIdsAndStatusAndIdGt(sessionIds,
  409 + Ai1SessionMessageStatusEnum.GENERATING.getStatus(), message1.getId()));
  410 + List<Ai1SessionMessageDO> limited = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds,
  411 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message1.getId(), null, 1);
  412 + List<Ai1SessionMessageDO> all = sessionMessageService.getSessionMessageListBySessionIdsAndStatusAndIdBetween(sessionIds,
  413 + Ai1SessionMessageStatusEnum.SUCCESS.getStatus(), message1.getId(), null, 100);
  414 + // 断言
  415 + assertEquals(List.of(message2.getId()), limited.stream().map(Ai1SessionMessageDO::getId).toList());
  416 + assertEquals(List.of(message2.getId(), message3.getId()), all.stream().map(Ai1SessionMessageDO::getId).toList());
  417 + }
  418 +
  419 + @Test
  420 + public void testGetGeneratingAssistantMessageIdList() {
  421 + // mock 数据:只有本会话生成中的助手消息命中
  422 + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole());
  423 + Ai1SessionMessageDO message1 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(),
  424 + Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  425 + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole());
  426 + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.FAILED.getStatus());
  427 + insertSessionMessage(10L, Ai1SessionMessageRoleEnum.USER.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  428 + Ai1SessionMessageDO message2 = insertSessionMessage(10L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(),
  429 + Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  430 + insertSessionMessage(11L, Ai1SessionMessageRoleEnum.ASSISTANT.getRole(), Ai1SessionMessageStatusEnum.GENERATING.getStatus());
  431 +
  432 + // 调用
  433 + List<Long> ids = sessionMessageService.getGeneratingAssistantMessageIdList(10L);
  434 + // 断言:只返回本会话生成中的助手消息编号,按编号升序
  435 + assertEquals(List.of(message1.getId(), message2.getId()), ids);
  436 + assertTrue(sessionMessageService.getGeneratingAssistantMessageIdList(12L).isEmpty());
  437 + }
  438 +
325 439 // ========== 测试数据 ==========
326 440  
327 441 private Ai1SessionMessageDO insertSessionMessage(Long sessionId, String role) {
  442 + return insertSessionMessage(sessionId, role, Ai1SessionMessageStatusEnum.SUCCESS.getStatus());
  443 + }
  444 +
  445 + private Ai1SessionMessageDO insertSessionMessage(Long sessionId, String role, Integer status) {
328 446 Ai1SessionMessageDO message = new Ai1SessionMessageDO().setSessionId(sessionId).setRole(role).setContent("内容")
329   - .setStatus(Ai1SessionMessageStatusEnum.SUCCESS.getStatus());
  447 + .setStatus(status);
330 448 sessionMessageMapper.insert(message);
331 449 return message;
332 450 }
... ...
yudao-module-ai1-springai/src/test/java/cn/iocoder/yudao/module/ai1/service/session/Ai1SessionServiceImplTest.java
1 1 package cn.iocoder.yudao.module.ai1.service.session;
2 2  
  3 +import cn.iocoder.yudao.framework.common.pojo.PageResult;
3 4 import cn.iocoder.yudao.framework.test.core.ut.BaseDbUnitTest;
4 5 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionCreateMyReqVO;
  6 +import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionPageReqVO;
5 7 import cn.iocoder.yudao.module.ai1.controller.admin.session.vo.session.Ai1SessionUpdateMyReqVO;
6 8 import cn.iocoder.yudao.module.ai1.dal.dataobject.session.Ai1SessionDO;
7 9 import cn.iocoder.yudao.module.ai1.dal.mysql.session.Ai1SessionMapper;
... ... @@ -104,4 +106,44 @@ public class Ai1SessionServiceImplTest extends BaseDbUnitTest {
104 106 assertEquals(Arrays.asList(id2, id1), list.stream().map(Ai1SessionDO::getId).toList());
105 107 }
106 108  
  109 + @Test
  110 + public void testGetUserIdListByAgentId() {
  111 + // mock 数据
  112 + sessionService.createSessionMy(200L, new Ai1SessionCreateMyReqVO().setAgentId(1L));
  113 + sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L));
  114 + sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L));
  115 + sessionService.createSessionMy(300L, new Ai1SessionCreateMyReqVO().setAgentId(2L));
  116 +
  117 + // 调用,并断言:去重并按用户编号升序
  118 + assertEquals(Arrays.asList(100L, 200L), sessionService.getUserIdListByAgentId(1L));
  119 + }
  120 +
  121 + @Test
  122 + public void testGetSession() {
  123 + // mock 数据
  124 + Long id = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L));
  125 +
  126 + // 调用,并断言
  127 + assertEquals(100L, sessionService.getSession(id).getUserId());
  128 + assertNull(sessionService.getSession(id + 1));
  129 + }
  130 +
  131 + @Test
  132 + public void testGetSessionPage_withoutAgentId() {
  133 + // mock 数据
  134 + Long id1 = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(1L));
  135 + Long id2 = sessionService.createSessionMy(100L, new Ai1SessionCreateMyReqVO().setAgentId(2L));
  136 + Long id3 = sessionService.createSessionMy(200L, new Ai1SessionCreateMyReqVO().setAgentId(2L));
  137 +
  138 + // 调用:不传 Agent 时跨 Agent 查询
  139 + PageResult<Ai1SessionDO> pageResult = sessionService.getSessionPage(new Ai1SessionPageReqVO());
  140 + // 断言
  141 + assertEquals(Arrays.asList(id3, id2, id1), pageResult.getList().stream().map(Ai1SessionDO::getId).toList());
  142 +
  143 + // 调用:按 Agent、用户筛选
  144 + pageResult = sessionService.getSessionPage(new Ai1SessionPageReqVO().setAgentId(2L).setUserId(100L));
  145 + // 断言
  146 + assertEquals(Collections.singletonList(id2), pageResult.getList().stream().map(Ai1SessionDO::getId).toList());
  147 + }
  148 +
107 149 }
... ...
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;
17 17  
18 18 public class Ai1UtilsTest {
19 19  
  20 + private static final List<String> MEMORY_SECTIONS = List.of("用户画像", "偏好习惯", "关键事实", "目标与意图", "近期进展");
  21 +
20 22 @Test
21 23 public void testNormalizeHeaders_orderAndValues() {
22 24 // 准备参数
... ... @@ -116,4 +118,55 @@ public class Ai1UtilsTest {
116 118 assertFalse(ex.getMessage().contains("literal-secret"));
117 119 }
118 120 }
  121 +
  122 + @Test
  123 + public void testNormalizeMemoryMarkdown_sections() {
  124 + // 调用,并断言:标题与内容压成一行时拆开,不同维度之间空一行
  125 + assertEquals("## 用户画像\n- 后端工程师\n\n## 偏好习惯\n- 简洁回答",
  126 + Ai1Utils.normalizeMemoryMarkdown("##用户画像-后端工程师##偏好习惯-简洁回答", MEMORY_SECTIONS));
  127 + // 调用,并断言:统一换行,合并多余空行
  128 + assertEquals("## 用户画像\n- 程序员\n\n## 关键事实\n- 住在北京",
  129 + Ai1Utils.normalizeMemoryMarkdown(" ## 用户画像\r\n-程序员\r\n\r\n\r\n## 关键事实\r- 住在北京 ", MEMORY_SECTIONS));
  130 + }
  131 +
  132 + @Test
  133 + public void testNormalizeMemoryMarkdown_keepsFacts() {
  134 + // 准备参数:C#、#123、URL 锚点、中文破折号都是用户事实,不能被拆开
  135 + String content = "## 关键事实\n- 主力语言 C#,工单#123,文档 https://a.cn/doc#part\n- 常驻北京—上海,周期 3–5 天";
  136 +
  137 + // 调用,并断言
  138 + assertEquals(content, Ai1Utils.normalizeMemoryMarkdown(content, MEMORY_SECTIONS));
  139 + // 调用,并断言:非固定维度的行内标题不拆行
  140 + assertEquals("## 关键事实\n- 标签##其他", Ai1Utils.normalizeMemoryMarkdown("## 关键事实\n- 标签##其他", MEMORY_SECTIONS));
  141 + }
  142 +
  143 + @Test
  144 + public void testNormalizeMemoryMarkdown_inlineList() {
  145 + // 调用,并断言:非 ASCII 字符后紧跟 ASCII「-」,拆成列表项
  146 + assertEquals("## 偏好习惯\n- 口吻正式\n- 开场先问候",
  147 + Ai1Utils.normalizeMemoryMarkdown("## 偏好习惯\n- 口吻正式-开场先问候", MEMORY_SECTIONS));
  148 + // 调用,并断言:ASCII 字符之间的「-」不拆
  149 + assertEquals("## 关键事实\n- 使用 spring-boot", Ai1Utils.normalizeMemoryMarkdown("## 关键事实\n- 使用 spring-boot", MEMORY_SECTIONS));
  150 + }
  151 +
  152 + @Test
  153 + public void testNormalizeMemoryMarkdown_blank() {
  154 + // 调用,并断言:为空时原样返回
  155 + assertNull(Ai1Utils.normalizeMemoryMarkdown(null, MEMORY_SECTIONS));
  156 + assertEquals(" ", Ai1Utils.normalizeMemoryMarkdown(" ", MEMORY_SECTIONS));
  157 + }
  158 +
  159 + @Test
  160 + public void testLimitLength() {
  161 + // 调用,并断言:先去除首尾空白,未超长时不截断
  162 + assertEquals("用户喜欢", Ai1Utils.limitLength(" 用户喜欢 \n", 4));
  163 + // 调用,并断言:下标 4、5 的 char 是同一个 emoji,截断不拆开代理对
  164 + assertEquals("用户喜欢", Ai1Utils.limitLength("用户喜欢😀表情", 5));
  165 + assertEquals("用户喜欢😀", Ai1Utils.limitLength("用户喜欢😀表情", 6));
  166 + assertEquals("用户喜欢😀表", Ai1Utils.limitLength("用户喜欢😀表情", 7));
  167 + // 调用,并断言:为空时原样返回
  168 + assertNull(Ai1Utils.limitLength(null, 5));
  169 + assertEquals(" ", Ai1Utils.limitLength(" ", 5));
  170 + }
  171 +
119 172 }
... ...
yudao-module-ai1-springai/src/test/resources/sql/clean.sql
... ... @@ -8,3 +8,4 @@ DELETE FROM &quot;ai1_knowledge_document&quot;;
8 8 DELETE FROM "ai1_agent";
9 9 DELETE FROM "ai1_session";
10 10 DELETE FROM "ai1_session_message";
  11 +DELETE FROM "ai1_memory";
... ...
yudao-module-ai1-springai/src/test/resources/sql/create_tables.sql
... ... @@ -152,6 +152,7 @@ CREATE TABLE IF NOT EXISTS &quot;ai1_agent&quot; (
152 152 "knowledge_base_ids" varchar(500) DEFAULT NULL,
153 153 "mcp_ids" varchar(500) DEFAULT NULL,
154 154 "skill_ids" varchar(500) DEFAULT NULL,
  155 + "memory_enabled" bit NOT NULL DEFAULT FALSE,
155 156 "status" tinyint NOT NULL DEFAULT 1,
156 157 "creator" varchar(64) DEFAULT '',
157 158 "create_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
... ... @@ -197,3 +198,21 @@ CREATE TABLE IF NOT EXISTS &quot;ai1_session_message&quot; (
197 198 "tenant_id" bigint NOT NULL DEFAULT 0,
198 199 PRIMARY KEY ("id")
199 200 );
  201 +
  202 +-- ----------------------------
  203 +-- AI1 长期记忆
  204 +-- ----------------------------
  205 +CREATE TABLE IF NOT EXISTS "ai1_memory" (
  206 + "id" bigint NOT NULL GENERATED BY DEFAULT AS IDENTITY,
  207 + "agent_id" bigint NOT NULL,
  208 + "user_id" bigint NOT NULL,
  209 + "content" mediumtext NOT NULL,
  210 + "last_message_id" bigint NOT NULL DEFAULT 0,
  211 + "creator" varchar(64) DEFAULT '',
  212 + "create_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
  213 + "updater" varchar(64) DEFAULT '',
  214 + "update_time" timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP,
  215 + "deleted" bit NOT NULL DEFAULT FALSE,
  216 + "tenant_id" bigint NOT NULL DEFAULT 0,
  217 + PRIMARY KEY ("id")
  218 +);
... ...