package vip.mate.channel.web;
import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import vip.mate.workspace.conversation.model.MessageContentPart;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
/**
* 聊天流状态追踪器
*
* 采用生产者-消费者解耦设计:将 SSE 事件的生产(Flux 订阅)与消费(SseEmitter 连接)解耦。
* 一个后台 Flux 生产者持续产出事件,广播给所有 SseEmitter 订阅者并缓存到 buffer。
* 新连接(重连)到来时,先回放 buffer,再接入实时流。
*
*
Single-instance assumption
* The {@link #runs} map is process-local memory. A reconnect
* request can only re-attach to a {@code RunState} that lives on the same
* JVM that originally created it. In a multi-node deployment behind a load
* balancer, the LB MUST be configured for sticky session by {@code conversationId}
* (Nginx {@code hash $arg_conversationId consistent;}, K8s Ingress
* cookie-based affinity, AWS ALB target-group stickiness, etc.).
*
*
This is an explicit CE constraint — see
* {@code rfcs/community/90-appendix/02-tech-debt-inventory.md §4.1} and
* {@code rfc-054 §0}. Cross-node SSE relay (Redis Stream / NATS / Kafka) is
* tracked under the EE roadmap.
*
*
Operator-facing diagnostics: callers should use
* {@link #streamExistsOnThisNode(String)} when distinguishing "stream finished
* normally" from "stream is on a different node" — both return {@code false}
* from {@link #attach(String, SseEmitter)} but mean very different things to
* the user.
*
* @author MateClaw Team
*/
@Slf4j
@Component
public class ChatStreamTracker {
/** buffer 最大事件数,超出后丢弃最早的 thinking_delta 事件以释放空间 */
private static final int MAX_BUFFER_SIZE = 16000;
private final ObjectMapper objectMapper;
/**
* Maximum size, in bytes, of a single SSE event JSON payload before
* {@link #broadcastChunked} splits the body into ordered
* {@code tool_result_chunk} events.
*/
static final int CHUNK_SIZE = 8192;
// ===== Configurable knobs (mateclaw.stream.*) =====
/**
* Gate for chunked tool-result transport. When {@code false},
* {@link #broadcastChunked} falls back to a single broadcast call so
* environments that prefer the legacy single-event behavior can opt out.
*/
@Value("${mateclaw.stream.chunked-tool-results:true}")
private boolean chunkedToolResultsEnabled = true;
/**
* Gate for {@code iteration_start} / {@code iteration_end} events emitted
* from graph nodes. Off-by-default deployments can suppress them without
* touching node code.
*/
@Value("${mateclaw.stream.iteration-events:true}")
private boolean iterationEventsEnabled = true;
/**
* Heartbeat cadence (seconds) before the first model token arrives. Short
* because pre-token gaps strand the UI on a blank "正在生成中" placeholder
* with no visible activity.
*/
@Value("${mateclaw.stream.heartbeat.pre-token-sec:2}")
private int heartbeatPreTokenSec = 2;
/**
* Heartbeat cadence (seconds) once the model is actively streaming tokens —
* deltas themselves keep the connection warm, so heartbeats relax.
*/
@Value("${mateclaw.stream.heartbeat.streaming-sec:10}")
private int heartbeatStreamingSec = 10;
/**
* Heartbeat cadence (seconds) while a tool call is in flight. Tools can
* take longer than streaming chunks but should still tick faster than the
* default proxy idle timeout.
*/
@Value("${mateclaw.stream.heartbeat.tool-sec:5}")
private int heartbeatToolSec = 5;
public ChatStreamTracker(ObjectMapper objectMapper) {
this.objectMapper = objectMapper;
}
/** Test-only setters; production paths use Spring property binding. */
void setChunkedToolResultsEnabled(boolean enabled) {
this.chunkedToolResultsEnabled = enabled;
}
void setIterationEventsEnabled(boolean enabled) {
this.iterationEventsEnabled = enabled;
}
public boolean isIterationEventsEnabled() {
return iterationEventsEnabled;
}
/**
* One buffered SSE event. The {@code id} is a per-conversation monotonic
* sequence — the SSE protocol's standard {@code id:} line carries this
* value so the client can echo it back via {@code lastEventId} when
* reconnecting, allowing us to skip already-delivered events on replay.
*/
record SseEvent(long id, String name, String json) {}
/**
* 中断类型:区分用户主动停止和用户在运行中追加新消息
*/
public enum InterruptType {
/** 用户点击 Stop,终止当前 turn,不自动续跑 */
USER_STOP,
/** 用户在执行中追加新消息,中断当前 turn 后自动续跑排队消息 */
USER_INTERRUPT_WITH_FOLLOWUP
}
static final class RunState {
final String conversationId;
final List subscribers = new ArrayList<>();
final List buffer = new ArrayList<>();
final Object lock = new Object();
volatile boolean done;
/**
* Monotonic sequence used as the SSE protocol {@code id:} field.
* Incremented inside {@code state.lock} as each event is buffered,
* so the buffer is always in (id-asc) order. On reconnect, the
* client echoes its last-seen id back via {@code lastEventId} and
* we skip events whose id is ≤ that value during replay —
* eliminating the duplicate-delivery class of bugs.
*/
long nextEventId = 0L;
/** Flux 订阅的 Disposable,用于取消 LLM 流 */
volatile Disposable disposable;
/** 停止标志:requestStop() 设为 true,各图节点和 LLM 调用检查此标志以提前退出 */
final AtomicBoolean stopRequested = new AtomicBoolean(false);
/**
* 当前活跃的 Flux 数量(原始流 + 审批 Replay 流共享同一个 RunState)。
* complete() 仅在计数归零时才真正移除 RunState,防止 Replay 仍在运行时被原始流的完成误删。
*/
volatile int activeFluxCount = 0;
// ===== Interrupt + Queue 新增字段 =====
/** 中断类型(null 表示未请求中断) */
volatile InterruptType interruptType;
/** 当前执行阶段(用于 heartbeat 和前端状态展示) */
volatile String currentPhase = "thinking";
/** 当前正在执行的工具名称 */
volatile String runningToolName;
/** 等待原因(审批等待时有值) */
volatile String waitingReason;
/** 排队的用户消息队列(支持多条排队消息,按序消费) */
final java.util.Queue messageQueue = new java.util.concurrent.ConcurrentLinkedQueue<>();
/**
* Emergency save callback registered by the SSE chain owner (ChatController).
* Invoked from {@link #onShutdown()} so the accumulated assistant content + tool_calls
* are persisted before the JVM tears down — without this, a `mvn spring-boot:run`
* restart wipes any in-flight turn and leaves only the user message in DB.
*
* The callback must be idempotent (will not be called twice for the same run, but
* may race with normal doOnComplete/doOnError; both paths must tolerate the other
* having saved already).
*/
volatile Runnable emergencySaveCallback;
/** 心跳定时器 */
volatile ScheduledFuture> heartbeatFuture;
/**
* Flips the first time any content/thinking delta is observed for this
* run. Heartbeat scheduling watches this flag to switch from the short
* pre-token cadence to the streaming cadence — pre-token gaps need
* frequent keep-alives because the UI has no other signal of activity.
*/
volatile boolean firstTokenReceived = false;
/** 已广播的 pending approval ID 集合(用于幂等去重) */
final java.util.Set broadcastedApprovalIds = java.util.concurrent.ConcurrentHashMap.newKeySet();
/** 创建时间(用于 stale 检测和清理) */
final long createdAt = System.currentTimeMillis();
/**
* Wall-clock millis of the most recent meaningful event on this run.
* Updated whenever {@link #broadcast(String, String, String)} pushes a
* non-heartbeat event so a watchdog can tell "actively producing"
* apart from "alive but silent".
*/
volatile long lastEventAt = System.currentTimeMillis();
/** Bound agent identifier; null while not yet resolved. */
volatile Long agentId;
/** Username that owns this run; null for system-driven runs. */
volatile String username;
RunState(String conversationId) {
this.conversationId = conversationId;
}
}
private final ConcurrentHashMap runs = new ConcurrentHashMap<>();
/**
* Conversations whose run was force-recycled by an admin. Maps to the
* recycle timestamp so a scheduled cleanup can age entries out (TTL
* matches {@link #DONE_RETENTION_MS} — long enough that any in-flight
* doOnComplete / doOnError firing after the dispose still finds the
* marker, short enough not to leak across sessions).
*
* Read by the SSE doOn* handlers in ChatController to skip a duplicate
* saveMessage when the recycle path already wrote the "[已被用户中止]"
* placeholder. Without this, the agent's late-yielding doOnComplete
* inserts a second assistant row carrying whatever the agent produced
* after the user pressed stop — exactly the behavior the user does
* not want when force-recycling.
*/
private final ConcurrentHashMap recycledConversations = new ConcurrentHashMap<>();
/** 事件 relay:子会话事件转发到父会话(用于 Agent 委派进度可见性) */
private final ConcurrentHashMap>> eventRelays = new ConcurrentHashMap<>();
/**
* 注册事件 relay:将 sourceConversationId 的广播事件同时转发给 listener。
* 返回一个 Runnable,调用后取消注册。
*/
public Runnable addEventRelay(String sourceConversationId,
java.util.function.BiConsumer listener) {
eventRelays.computeIfAbsent(sourceConversationId, k -> new java.util.concurrent.CopyOnWriteArrayList<>())
.add(listener);
log.debug("Event relay registered for conversation {}", sourceConversationId);
return () -> {
List> listeners = eventRelays.get(sourceConversationId);
if (listeners != null) {
listeners.remove(listener);
if (listeners.isEmpty()) {
eventRelays.remove(sourceConversationId);
}
}
log.debug("Event relay removed for conversation {}", sourceConversationId);
};
}
/**
* Batching variant of {@link #addEventRelay} for sub-conversation streams
* whose tool-call chatter would flood the parent transcript. Tool start /
* complete events accumulate into a buffer; lifecycle and error events
* (subagent_*, error, tool_approval_requested, phase, done) bypass the
* buffer but flush it first so ordering is preserved.
*
* Buffered events are emitted as a single {@code delegation_batch}
* envelope on the parent conversation listener:
*
* {
* "kind": "delegation_batch",
* "scope": "subagent",
* "events": [{ "event": "tool_call_started", "data": "<json>" }, ...]
* }
*
*
* @param sourceConversationId conversation to listen on
* @param parentConversationId parent conversation context (currently
* forwarded only as listener metadata; the
* tracker itself does not target it)
* @param batchSize flush threshold by event count
* @param flushMs flush threshold by elapsed millis since
* first buffered event
* @return Runnable that deregisters the relay (and flushes any pending
* events first)
*/
public Runnable addBatchedEventRelay(String sourceConversationId,
String parentConversationId,
int batchSize,
long flushMs,
java.util.function.BiConsumer listener) {
BatchedRelay relay = new BatchedRelay(parentConversationId, listener,
Math.max(1, batchSize), Math.max(1, flushMs));
java.util.function.BiConsumer wrapper = relay::accept;
eventRelays.computeIfAbsent(sourceConversationId,
k -> new java.util.concurrent.CopyOnWriteArrayList<>())
.add(wrapper);
log.debug("Batched relay registered for conversation {} -> parent={}",
sourceConversationId, parentConversationId);
return () -> {
relay.shutdown();
List> listeners =
eventRelays.get(sourceConversationId);
if (listeners != null) {
listeners.remove(wrapper);
if (listeners.isEmpty()) {
eventRelays.remove(sourceConversationId);
}
}
log.debug("Batched relay removed for {} -> parent={}",
sourceConversationId, parentConversationId);
};
}
/**
* Internal helper holding the batch buffer and the scheduled flush. Each
* relay owns its own state but reuses {@link #heartbeatScheduler} for
* timer ticks (sharing the daemon-thread scheduler avoids one-thread-per
* -relay sprawl in long agent sessions).
*/
private final class BatchedRelay {
private final String parentConversationId;
private final java.util.function.BiConsumer downstream;
private final int batchSize;
private final long flushMs;
private final List