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
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("*** [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<Message> 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()) {
@@ -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);
}
}
}
@@ -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();
@@ -3,18 +3,17 @@ package com.superbiz.agent.util;
/**
* 会话上下文持有者(基于 ThreadLocal)
* <p>
* 用于在执行链路中传递 sessionId 和 agentName,覆盖 AgentLoggingHook 和
* LookupKnowledgeTool 等无法直接通过 RunnableConfig 获取上下文的组件。
* 用于在同步调用链路中传递 sessionId,兜底 LookupKnowledgeTool 等
* 无法通过 RunnableConfig 获取上下文的组件。
* 优先使用 RunnableConfig.metadata 传递,ThreadLocal 作为同步路径的补充。
* <p>
* 使用规范:
* 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<String> SESSION_ID = new ThreadLocal<>();
private static final ThreadLocal<String> 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();
}
}