From a3abe3f7a2e9e27b7a89596204f256028697eb5c Mon Sep 17 00:00:00 2001 From: zhuyongxin Date: Fri, 26 Jun 2026 17:28:30 +0800 Subject: [PATCH] =?UTF-8?q?refactor(session):=20=E6=B8=85=E7=90=86?= =?UTF-8?q?=E4=BB=A3=E7=A0=81=20&=20RunnableConfig=20=E4=BC=A0=20sessionId?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - AgentLoggingHook 改为从 config.metadata 读取 sessionId（线程安全） - 移除 AgentLoggingHook 调试用的 metadata 日志 - TokenTrackingChatModel 日志降为 debug - SessionContextHolder 移除未使用的 setAgentName/getAgentName - ChatService 清理无用 import - 修复 stream 路径下 ThreadLocal NPE --- .../superbiz/agent/hook/AgentLoggingHook.java | 23 +++++++++++-------- .../agent/hook/TokenTrackingChatModel.java | 19 +++++++++------ .../superbiz/agent/service/ChatService.java | 8 ++++++- .../agent/util/SessionContextHolder.java | 20 ++++------------ 4 files changed, 37 insertions(+), 33 deletions(-) diff --git a/src/main/java/com/superbiz/agent/hook/AgentLoggingHook.java b/src/main/java/com/superbiz/agent/hook/AgentLoggingHook.java index 5f79d99..b6ec4dd 100644 --- a/src/main/java/com/superbiz/agent/hook/AgentLoggingHook.java +++ b/src/main/java/com/superbiz/agent/hook/AgentLoggingHook.java @@ -46,12 +46,20 @@ public class AgentLoggingHook extends MessagesModelHook { @Override public AgentCommand beforeModel(List previousMessages, RunnableConfig config) { - String sessionId = SessionContextHolder.getSessionId(); + // 优先从 config.metadata 取 sessionId（线程安全），兜底 ThreadLocal + String sessionId = config.metadata("sessionId") + .map(Object::toString) + .orElseGet(SessionContextHolder::getSessionId); - int stepIndex = stepCounters.merge(sessionId, 0, (old, one) -> old + 1); + boolean hasSession = (sessionId != null); + + int stepIndex = 0; + if (hasSession) { + stepIndex = stepCounters.merge(sessionId, 0, (old, one) -> old + 1); + } log.info("========================================"); - log.info("*** [Agent 思考] 第 {} 轮思考开始", stepIndex + 1); + log.info("*** [Agent 思考] 第 {} 轮思考开始", (hasSession ? stepCounters.get(sessionId) : 0) + 1); log.info("*** [Agent 思考] 当前消息数量: {}", previousMessages.size()); // 打印最后几条消息 @@ -101,9 +109,10 @@ public class AgentLoggingHook extends MessagesModelHook { @Override public AgentCommand afterModel(List previousMessages, RunnableConfig config) { String sessionId = SessionContextHolder.getSessionId(); + boolean hasSession = (sessionId != null); log.info("========================================"); - log.info("*** [Agent 思考] 第 {} 轮思考完成", stepCounters.getOrDefault(sessionId, 0)); + log.info("*** [Agent 思考] 第 {} 轮思考完成", (hasSession ? stepCounters.getOrDefault(sessionId, 0) : 0)); // 查找最后一条 AssistantMessage（模型的回复） AssistantMessage lastAssistant = null; @@ -117,12 +126,6 @@ public class AgentLoggingHook extends MessagesModelHook { boolean hasToolCall = false; if (lastAssistant != null) { - // 调试：打印 metadata - if (lastAssistant.getMetadata() != null && !lastAssistant.getMetadata().isEmpty()) { - log.info("*** [Agent 思考] 模型返回 metadata: {}", lastAssistant.getMetadata()); - } else { - log.info("*** [Agent 思考] 模型返回 metadata: (空)"); - } // 打印模型返回的文本内容 String textContent = extractTextContent(lastAssistant); if (textContent != null && !textContent.isEmpty()) { diff --git a/src/main/java/com/superbiz/agent/hook/TokenTrackingChatModel.java b/src/main/java/com/superbiz/agent/hook/TokenTrackingChatModel.java index 92af340..f9adc02 100644 --- a/src/main/java/com/superbiz/agent/hook/TokenTrackingChatModel.java +++ b/src/main/java/com/superbiz/agent/hook/TokenTrackingChatModel.java @@ -1,5 +1,7 @@ package com.superbiz.agent.hook; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.prompt.Prompt; @@ -11,6 +13,8 @@ import reactor.core.publisher.Flux; */ public class TokenTrackingChatModel implements ChatModel { + private static final Logger log = LoggerFactory.getLogger(TokenTrackingChatModel.class); + private final ChatModel delegate; public TokenTrackingChatModel(ChatModel delegate) { @@ -31,15 +35,16 @@ public class TokenTrackingChatModel implements ChatModel { private void captureTokenUsage(ChatResponse response) { try { - if (response.getMetadata() != null && response.getMetadata().getUsage() != null) { - var usage = response.getMetadata().getUsage(); - Integer total = usage.getTotalTokens(); - if (total != null && total > 0) { - TokenUsageHolder.set(total); - } + if (response.getMetadata() == null || response.getMetadata().getUsage() == null) { + return; + } + var usage = response.getMetadata().getUsage(); + Integer total = usage.getTotalTokens(); + if (total != null && total > 0) { + TokenUsageHolder.set(total); } } catch (Exception e) { - // 不中断模型调用 + log.debug("捕获 token 用量失败", e); } } } diff --git a/src/main/java/com/superbiz/agent/service/ChatService.java b/src/main/java/com/superbiz/agent/service/ChatService.java index 8ba483d..c311fd7 100644 --- a/src/main/java/com/superbiz/agent/service/ChatService.java +++ b/src/main/java/com/superbiz/agent/service/ChatService.java @@ -1,6 +1,7 @@ package com.superbiz.agent.service; import com.alibaba.cloud.ai.graph.OverAllState; +import com.alibaba.cloud.ai.graph.RunnableConfig; import com.alibaba.cloud.ai.graph.agent.ReactAgent; import com.alibaba.cloud.ai.graph.agent.flow.agent.SupervisorAgent; import com.alibaba.cloud.ai.graph.exception.GraphRunnerException; @@ -254,7 +255,12 @@ public class ChatService { SessionContextHolder.setSessionId(sessionId); try { - var response = agent.call(question); + // 通过 RunnableConfig 将 sessionId 传入 Hook（线程安全，异步也兼容） + var config = RunnableConfig.builder() + .addMetadata("sessionId", sessionId) + .build(); + + var response = agent.call(question, config); long duration = System.currentTimeMillis() - startTime; String answer = response.getText(); diff --git a/src/main/java/com/superbiz/agent/util/SessionContextHolder.java b/src/main/java/com/superbiz/agent/util/SessionContextHolder.java index 58bfd37..f0e0f7a 100644 --- a/src/main/java/com/superbiz/agent/util/SessionContextHolder.java +++ b/src/main/java/com/superbiz/agent/util/SessionContextHolder.java @@ -3,18 +3,17 @@ package com.superbiz.agent.util; /** * 会话上下文持有者（基于 ThreadLocal） *

- * 用于在执行链路中传递 sessionId 和 agentName，覆盖 AgentLoggingHook 和 - * LookupKnowledgeTool 等无法直接通过 RunnableConfig 获取上下文的组件。 + * 用于在同步调用链路中传递 sessionId，兜底 LookupKnowledgeTool 等 + * 无法通过 RunnableConfig 获取上下文的组件。 + * 优先使用 RunnableConfig.metadata 传递，ThreadLocal 作为同步路径的补充。 *

* 使用规范： - * 1. 调用方（ChatService/AiOpsService）在 Agent 执行前调用 setSessionId() 和 setAgentName() - * 2. AgentLoggingHook 和工具类通过 getSessionId() / getAgentName() 读取 - * 3. 必须在 finally 块中调用 clear()，防止内存泄漏和线程污染 + * 1. 调用方在 Agent 执行前调用 setSessionId() + * 2. finally 块中调用 clear() */ public class SessionContextHolder { private static final ThreadLocal SESSION_ID = new ThreadLocal<>(); - private static final ThreadLocal AGENT_NAME = new ThreadLocal<>(); public static void setSessionId(String sessionId) { SESSION_ID.set(sessionId); @@ -24,16 +23,7 @@ public class SessionContextHolder { return SESSION_ID.get(); } - public static void setAgentName(String agentName) { - AGENT_NAME.set(agentName); - } - - public static String getAgentName() { - return AGENT_NAME.get(); - } - public static void clear() { SESSION_ID.remove(); - AGENT_NAME.remove(); } }