refactor(session): 清理代码 & RunnableConfig 传 sessionId

- AgentLoggingHook 改为从 config.metadata 读取 sessionId(线程安全)
- 移除 AgentLoggingHook 调试用的 metadata 日志
- TokenTrackingChatModel 日志降为 debug
- SessionContextHolder 移除未使用的 setAgentName/getAgentName
- ChatService 清理无用 import
- 修复 stream 路径下 ThreadLocal NPE
This commit is contained in:
zhuyongxin
2026-06-26 17:28:30 +08:00
parent 0d9cce75f9
commit a3abe3f7a2
4 changed files with 37 additions and 33 deletions
@@ -46,12 +46,20 @@ public class AgentLoggingHook extends MessagesModelHook {
@Override @Override
public AgentCommand beforeModel(List<Message> previousMessages, RunnableConfig config) { public AgentCommand beforeModel(List<Message> 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("========================================");
log.info("*** [Agent 思考] 第 {} 轮思考开始", stepIndex + 1); log.info("*** [Agent 思考] 第 {} 轮思考开始", (hasSession ? stepCounters.get(sessionId) : 0) + 1);
log.info("*** [Agent 思考] 当前消息数量: {}", previousMessages.size()); log.info("*** [Agent 思考] 当前消息数量: {}", previousMessages.size());
// 打印最后几条消息 // 打印最后几条消息
@@ -101,9 +109,10 @@ public class AgentLoggingHook extends MessagesModelHook {
@Override @Override
public AgentCommand afterModel(List<Message> previousMessages, RunnableConfig config) { public AgentCommand afterModel(List<Message> previousMessages, RunnableConfig config) {
String sessionId = SessionContextHolder.getSessionId(); String sessionId = SessionContextHolder.getSessionId();
boolean hasSession = (sessionId != null);
log.info("========================================"); log.info("========================================");
log.info("*** [Agent 思考] 第 {} 轮思考完成", stepCounters.getOrDefault(sessionId, 0)); log.info("*** [Agent 思考] 第 {} 轮思考完成", (hasSession ? stepCounters.getOrDefault(sessionId, 0) : 0));
// 查找最后一条 AssistantMessage(模型的回复) // 查找最后一条 AssistantMessage(模型的回复)
AssistantMessage lastAssistant = null; AssistantMessage lastAssistant = null;
@@ -117,12 +126,6 @@ public class AgentLoggingHook extends MessagesModelHook {
boolean hasToolCall = false; boolean hasToolCall = false;
if (lastAssistant != null) { 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); String textContent = extractTextContent(lastAssistant);
if (textContent != null && !textContent.isEmpty()) { if (textContent != null && !textContent.isEmpty()) {
@@ -1,5 +1,7 @@
package com.superbiz.agent.hook; 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.ChatModel;
import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.chat.prompt.Prompt;
@@ -11,6 +13,8 @@ import reactor.core.publisher.Flux;
*/ */
public class TokenTrackingChatModel implements ChatModel { public class TokenTrackingChatModel implements ChatModel {
private static final Logger log = LoggerFactory.getLogger(TokenTrackingChatModel.class);
private final ChatModel delegate; private final ChatModel delegate;
public TokenTrackingChatModel(ChatModel delegate) { public TokenTrackingChatModel(ChatModel delegate) {
@@ -31,15 +35,16 @@ public class TokenTrackingChatModel implements ChatModel {
private void captureTokenUsage(ChatResponse response) { private void captureTokenUsage(ChatResponse response) {
try { try {
if (response.getMetadata() != null && response.getMetadata().getUsage() != null) { if (response.getMetadata() == null || response.getMetadata().getUsage() == null) {
var usage = response.getMetadata().getUsage(); return;
Integer total = usage.getTotalTokens(); }
if (total != null && total > 0) { var usage = response.getMetadata().getUsage();
TokenUsageHolder.set(total); Integer total = usage.getTotalTokens();
} if (total != null && total > 0) {
TokenUsageHolder.set(total);
} }
} catch (Exception e) { } catch (Exception e) {
// 不中断模型调用 log.debug("捕获 token 用量失败", e);
} }
} }
} }
@@ -1,6 +1,7 @@
package com.superbiz.agent.service; package com.superbiz.agent.service;
import com.alibaba.cloud.ai.graph.OverAllState; 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.ReactAgent;
import com.alibaba.cloud.ai.graph.agent.flow.agent.SupervisorAgent; import com.alibaba.cloud.ai.graph.agent.flow.agent.SupervisorAgent;
import com.alibaba.cloud.ai.graph.exception.GraphRunnerException; import com.alibaba.cloud.ai.graph.exception.GraphRunnerException;
@@ -254,7 +255,12 @@ public class ChatService {
SessionContextHolder.setSessionId(sessionId); SessionContextHolder.setSessionId(sessionId);
try { 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; long duration = System.currentTimeMillis() - startTime;
String answer = response.getText(); String answer = response.getText();
@@ -3,18 +3,17 @@ package com.superbiz.agent.util;
/** /**
* 会话上下文持有者(基于 ThreadLocal) * 会话上下文持有者(基于 ThreadLocal)
* <p> * <p>
* 用于在执行链路中传递 sessionId 和 agentName,覆盖 AgentLoggingHook 和 * 用于在同步调用链路中传递 sessionId,兜底 LookupKnowledgeTool 等
* LookupKnowledgeTool 等无法直接通过 RunnableConfig 获取上下文的组件。 * 无法通过 RunnableConfig 获取上下文的组件。
* 优先使用 RunnableConfig.metadata 传递,ThreadLocal 作为同步路径的补充。
* <p> * <p>
* 使用规范: * 使用规范:
* 1. 调用方(ChatService/AiOpsService)在 Agent 执行前调用 setSessionId() 和 setAgentName() * 1. 调用方在 Agent 执行前调用 setSessionId()
* 2. AgentLoggingHook 和工具类通过 getSessionId() / getAgentName() 读取 * 2. finally 块中调用 clear()
* 3. 必须在 finally 块中调用 clear(),防止内存泄漏和线程污染
*/ */
public class SessionContextHolder { public class SessionContextHolder {
private static final ThreadLocal<String> SESSION_ID = new ThreadLocal<>(); private static final ThreadLocal<String> SESSION_ID = new ThreadLocal<>();
private static final ThreadLocal<String> AGENT_NAME = new ThreadLocal<>();
public static void setSessionId(String sessionId) { public static void setSessionId(String sessionId) {
SESSION_ID.set(sessionId); SESSION_ID.set(sessionId);
@@ -24,16 +23,7 @@ public class SessionContextHolder {
return SESSION_ID.get(); 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() { public static void clear() {
SESSION_ID.remove(); SESSION_ID.remove();
AGENT_NAME.remove();
} }
} }