diff --git a/mateclaw-server/src/main/java/vip/mate/agent/AgentService.java b/mateclaw-server/src/main/java/vip/mate/agent/AgentService.java index cd92cf83..3b4c886f 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentService.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentService.java @@ -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. + * + *

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 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 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 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); + } + } } diff --git a/mateclaw-server/src/main/java/vip/mate/channel/ChannelMessageRouter.java b/mateclaw-server/src/main/java/vip/mate/channel/ChannelMessageRouter.java index f3bbbeef..4a9c5bc6 100644 --- a/mateclaw-server/src/main/java/vip/mate/channel/ChannelMessageRouter.java +++ b/mateclaw-server/src/main/java/vip/mate/channel/ChannelMessageRouter.java @@ -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); diff --git a/mateclaw-server/src/main/java/vip/mate/channel/web/ChatController.java b/mateclaw-server/src/main/java/vip/mate/channel/web/ChatController.java index f6d453b8..270bbf00 100644 --- a/mateclaw-server/src/main/java/vip/mate/channel/web/ChatController.java +++ b/mateclaw-server/src/main/java/vip/mate/channel/web/ChatController.java @@ -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); } diff --git a/mateclaw-server/src/main/java/vip/mate/channel/web/TalkModeWebSocketHandler.java b/mateclaw-server/src/main/java/vip/mate/channel/web/TalkModeWebSocketHandler.java index 45bd703c..89cea694 100644 --- a/mateclaw-server/src/main/java/vip/mate/channel/web/TalkModeWebSocketHandler.java +++ b/mateclaw-server/src/main/java/vip/mate/channel/web/TalkModeWebSocketHandler.java @@ -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, diff --git a/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobLifecycleService.java b/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobLifecycleService.java index fcb01f50..cfc7f123 100644 --- a/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobLifecycleService.java +++ b/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobLifecycleService.java @@ -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 diff --git a/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobRunner.java b/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobRunner.java index 9b3f453d..d8273750 100644 --- a/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobRunner.java +++ b/mateclaw-server/src/main/java/vip/mate/cron/service/CronJobRunner.java @@ -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); } /**