opensource/mateclaw
2
0
Fork 0
mirror of https://gitee.com/mateos/mateclaw.git synced 2026-09-13 03:13:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat(conversation): introduce ChatResult to carry token usage through sync chat paths

Browse Source
This commit is contained in:
matevip 2026-05-25 21:15:53 +08:00
parent 8d01396130
commit a01f0354eb
6 changed files with 127 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

76
mateclaw-server/src/main/java/vip/mate/agent/AgentService.java
View File

@ -25,6 +25,7 @@ import vip.mate.workspace.conversation.model.ConversationEntity;
import vip.mate.workspace.conversation.repository.ConversationMapper;
import java.util.List;
import java.time.Duration;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;
@ -237,6 +238,26 @@ public class AgentService {
}
}
/**
* Sync chat that also captures token usage and runtime model attribution
* from the agent graph's {@code _usage_final} event. Equivalent to
* subscribing to {@link #chatStructuredStream} and joining all content
* deltas — produces the same assistant text as {@link #chat} but exposes
* the usage figures so callers can persist them on the assistant message.
*
* <p>Prefer this entry over {@link #chat} for any path that writes the
* reply to {@code mate_message} (sync HTTP endpoint, voice WebSocket,
* cron task, post-approval replay); the plain {@link #chat} stays as the
* thin wrapper for fire-and-forget invocations where usage is not needed.
*/
public ChatResult chatWithUsage(Long agentId, String message, String conversationId) {
return chatWithUsage(agentId, message, conversationId, ChatOrigin.EMPTY);
}
public ChatResult chatWithUsage(Long agentId, String message, String conversationId, ChatOrigin origin) {
return collectChatResult(chatStructuredStream(agentId, message, conversationId, "", null, origin));
}
public Flux<String> chatStream(Long agentId, String message, String conversationId) {
return chatStream(agentId, message, conversationId, ChatOrigin.EMPTY);
}
@ -362,6 +383,42 @@ public class AgentService {
}
}
/**
* Replay-after-approval that also captures token usage and runtime model
* attribution. Mirrors {@link #chatWithUsage} for the
* approval-resumption path used by {@code ChannelMessageRouter}.
*/
public ChatResult chatWithReplayWithUsage(Long agentId, String userMessage, String conversationId,
String toolCallPayload, ChatOrigin origin) {
return collectChatResult(chatWithReplayStream(agentId, userMessage, conversationId,
toolCallPayload, "", origin != null ? origin : ChatOrigin.EMPTY));
}
/**
* Subscribe to a structured stream and collapse it into a single
* {@link ChatResult}: append all content deltas, capture the trailing
* {@code _usage_final} event for token and model attribution.
*/
private ChatResult collectChatResult(Flux<StreamDelta> stream) {
StringBuilder content = new StringBuilder();
final int[] usage = {0, 0};
final String[] modelInfo = {null, null};
stream.doOnNext(delta -> {
if (delta.isEvent() && "_usage_final".equals(delta.eventType())) {
Map<String, Object> data = delta.eventData();
usage[0] = ((Number) data.getOrDefault("promptTokens", 0)).intValue();
usage[1] = ((Number) data.getOrDefault("completionTokens", 0)).intValue();
Object model = data.get("runtimeModelName");
Object provider = data.get("runtimeProviderId");
if (model != null) modelInfo[0] = model.toString();
if (provider != null) modelInfo[1] = provider.toString();
} else if (delta.content() != null) {
content.append(delta.content());
}
}).blockLast(Duration.ofMinutes(10));
return new ChatResult(content.toString(), usage[0], usage[1], modelInfo[0], modelInfo[1]);
}
/**
* 带工具重放的流式调用（Web 端审批通过后使用，通过 SSE 推送结果）
*/
@ -623,4 +680,23 @@ public class AgentService {
return thinking != null ? thinking.length() : 0;
}
}
// ==================== ChatResult ====================
/**
* Sync chat result carrying the assistant reply alongside the usage
* attribution that the streaming path exposes via the {@code _usage_final}
* event. Use this when callers need to persist {@code promptTokens} /
* {@code completionTokens} / {@code runtimeModel} / {@code runtimeProvider}
* on the assistant message row but cannot subscribe to the structured
* stream directly (cron tasks, sync HTTP endpoints, voice WebSocket,
* post-approval replays).
*/
public record ChatResult(String content, int promptTokens, int completionTokens,
String runtimeModel, String runtimeProvider) {
public static ChatResult contentOnly(String content) {
return new ChatResult(content != null ? content : "", 0, 0, null, null);
}
}
}

7
mateclaw-server/src/main/java/vip/mate/channel/ChannelMessageRouter.java
View File

@ -1017,8 +1017,9 @@ public class ChannelMessageRouter {
replayOrigin = chatOriginFactory.from(
channelEntity, triggerMessage, conversationId, /* workspaceBasePath */ null);
}
String reply = agentService.chatWithReplay(
AgentService.ChatResult replayResult = agentService.chatWithReplayWithUsage(
agentId, replayPrompt, conversationId, consumed.getToolCallPayload(), replayOrigin);
String reply = replayResult.content();
// Persist the replay result. If the LLM 400'd during replay,
// the error reply must also get status='error' — otherwise the
@ -1026,7 +1027,9 @@ public class ChannelMessageRouter {
// into the prompt and re-trigger the same failure.
boolean isError = errorClassifier.isErrorReply(reply);
conversationService.saveMessage(conversationId, "assistant", reply, null,
isError ? "error" : "completed");
isError ? "error" : "completed",
replayResult.promptTokens(), replayResult.completionTokens(),
replayResult.runtimeModel(), replayResult.runtimeProvider());
// 发送回复
adapter.renderAndSend(replyTarget, reply);

7
mateclaw-server/src/main/java/vip/mate/channel/web/ChatController.java
View File

@ -1035,8 +1035,11 @@ public class ChatController {
conversationService.saveMessage(request.getConversationId(), "user", request.getMessage(), request.getContentParts());
String promptText = buildPromptText(request.getMessage(), request.getContentParts());
String response = agentService.chat(agentId, promptText, request.getConversationId());
conversationService.saveMessage(request.getConversationId(), "assistant", response);
AgentService.ChatResult result = agentService.chatWithUsage(agentId, promptText, request.getConversationId());
String response = result.content();
conversationService.saveMessage(request.getConversationId(), "assistant", response, null, "completed",
result.promptTokens(), result.completionTokens(),
result.runtimeModel(), result.runtimeProvider());
completionPublisher.publish(agentId, request.getConversationId(), request.getMessage(), response, "web");
return R.ok(response);
}

10
mateclaw-server/src/main/java/vip/mate/channel/web/TalkModeWebSocketHandler.java
View File

@ -145,13 +145,17 @@ public class TalkModeWebSocketHandler extends AbstractWebSocketHandler {
conversationService.saveMessage(talkSession.conversationId, "user", transcript, List.of());
// 5. Agent 对话（同步）
String reply = agentService.chat(talkSession.agentId, transcript, talkSession.conversationId);
AgentService.ChatResult chatResult = agentService.chatWithUsage(
talkSession.agentId, transcript, talkSession.conversationId);
String reply = chatResult.content();
if (reply == null || reply.isBlank()) {
reply = "Sorry, I couldn't generate a response.";
}
// 6. 保存助手回复
conversationService.saveMessage(talkSession.conversationId, "assistant", reply, List.of());
// 6. 保存助手回复（携带 token usage + runtime model 归属）
conversationService.saveMessage(talkSession.conversationId, "assistant", reply, List.of(),
"completed", chatResult.promptTokens(), chatResult.completionTokens(),
chatResult.runtimeModel(), chatResult.runtimeProvider());
// Publish conversation-completed event so memory extraction runs for voice turns too.
completionPublisher.publish(talkSession.agentId, talkSession.conversationId,

23
mateclaw-server/src/main/java/vip/mate/cron/service/CronJobLifecycleService.java
View File

@ -9,6 +9,7 @@ import org.springframework.context.ApplicationEventPublisher;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Propagation;
import org.springframework.transaction.annotation.Transactional;
import vip.mate.agent.AgentService;
import vip.mate.cron.delivery.CronJobCompletedEvent;
import vip.mate.cron.model.CronJobEntity;
import vip.mate.dashboard.model.CronJobRunEntity;
@ -178,6 +179,20 @@ public class CronJobLifecycleService {
public void finishRunAndPublish(CronJobEntity job, CronJobRunEntity run,
String userMessage, AssistantMessage result,
String conversationId, boolean silent) {
finishRunAndPublish(job, run, userMessage, result, conversationId, silent, null);
}
/**
* @param chatResult optional usage attribution from the LLM path; pass
* {@code null} for non-LLM paths (e.g. reminder
* direct-push) so the assistant row is persisted with
* zero token counts and null runtime model attribution.
*/
@Transactional(propagation = Propagation.REQUIRES_NEW)
public void finishRunAndPublish(CronJobEntity job, CronJobRunEntity run,
String userMessage, AssistantMessage result,
String conversationId, boolean silent,
AgentService.ChatResult chatResult) {
String convId = conversationId != null ? conversationId : run.getConversationId();
String text = result != null && result.getText() != null ? result.getText() : "";
@ -197,7 +212,13 @@ public class CronJobLifecycleService {
return;
}
conversationService.saveMessage(convId, "assistant", text);
if (chatResult != null) {
conversationService.saveMessage(convId, "assistant", text, null, "completed",
chatResult.promptTokens(), chatResult.completionTokens(),
chatResult.runtimeModel(), chatResult.runtimeProvider());
} else {
conversationService.saveMessage(convId, "assistant", text);
}
// Memory pipeline (existing behavior preserved — was inline in the
// old executeJob; now lives behind the same publisher used by the

21
mateclaw-server/src/main/java/vip/mate/cron/service/CronJobRunner.java
View File

@ -138,10 +138,12 @@ public class CronJobRunner {
// No-tx segment — long LLM call. RFC §5.2 hard rule: must not hold
// any DB connection during this call.
AgentService.ChatResult chatResult;
AssistantMessage result;
try {
ChatOrigin origin = originFactory.from(job, conversationId);
result = runAgent(job, userMessage, origin, conversationId);
chatResult = runAgent(job, userMessage, origin, conversationId);
result = new AssistantMessage(chatResult.content());
} catch (Exception e) {
log.error("[CronRunner] runAgent failed for job {}: {}", job.getId(), e.getMessage(), e);
try {
@ -156,12 +158,12 @@ public class CronJobRunner {
// Explicit no-op: the agent answered with the silent sentinel,
// meaning there is nothing to deliver or report for this run.
boolean silent = result != null && result.getText() != null
boolean silent = result.getText() != null
&& CRON_SILENT_MARKER.equals(result.getText().trim());
// T2 — short tx
try {
lifecycle.finishRunAndPublish(job, run, userMessage, result, conversationId, silent);
lifecycle.finishRunAndPublish(job, run, userMessage, result, conversationId, silent, chatResult);
} catch (Exception e) {
log.error("[CronRunner] T2 finishRunAndPublish failed for job {}: {}", job.getId(), e.getMessage(), e);
try {
@ -261,13 +263,14 @@ public class CronJobRunner {
* Runs the agent with the scheduled-job {@link ChatOrigin} and the
* execution-context prompt assembled by {@link #buildCronPrompt}.
*/
private AssistantMessage runAgent(CronJobEntity job, String userMessage, ChatOrigin origin,
String conversationId) {
private AgentService.ChatResult runAgent(CronJobEntity job, String userMessage, ChatOrigin origin,
String conversationId) {
String prompt = buildCronPrompt(userMessage, origin);
String text = "agent".equals(job.getTaskType())
? agentService.execute(job.getAgentId(), prompt, conversationId, origin)
: agentService.chat(job.getAgentId(), prompt, conversationId, origin);
return new AssistantMessage(text != null ? text : "");
// execute() and chat() both ultimately route through the agent's
// StateGraph; chatWithUsage captures token + runtime model attribution
// for either path. Plan-Execute agents stream via the same
// chatStructuredStream the helper consumes.
return agentService.chatWithUsage(job.getAgentId(), prompt, conversationId, origin);
}
/**