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

fix(webchat): harden stream cleanup and timeout handling

Browse Source 
Normalize stream idle timeout overrides, make WebChat orphan cleanup use exact run handles so stale callbacks cannot mutate replacement runs, and preserve approval replay usage metadata with regression coverage.
This commit is contained in:
倪程伟 2026-08-12 14:08:59 +08:00 committed by GitHub
parent 438d8d26ae
commit 0ae8552f5a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
10 changed files with 1167 additions and 228 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java
View File

@ -36,6 +36,7 @@ import org.springframework.beans.factory.annotation.Autowired;
import vip.mate.config.GraphObservationProperties;
import vip.mate.config.ReasoningRetentionProperties;
import vip.mate.exception.MateClawException;
import vip.mate.llm.chatmodel.HttpTimeouts;
import vip.mate.llm.chatmodel.OpenAiCompatibleChatModelBuilder;
import vip.mate.llm.chatmodel.ReasoningEffortResolver;
import vip.mate.llm.model.ModelConfigEntity;
@ -660,16 +661,9 @@ public class AgentGraphBuilder {
contextWindowResolver.noteContextLimitError(
primaryModelConfig.getProvider(),
primaryModelConfig.getModelName(), errorMessage));
// Issue #585: drive the streaming inter-frame idle timeout from
// the per-model read-timeout knob so a stalled provider can't
// hang the body Flux after the response headers arrive. Only
// override when the model explicitly sets a value — otherwise
// the helper keeps its 180s default.
Integer perModelTimeout = primaryModelConfig.getRequestTimeoutSeconds();
if (perModelTimeout != null) {
streamingHelper.setStreamIdleTimeoutSec(perModelTimeout);
}
}
streamingHelper.setStreamIdleTimeoutSec(
resolveStreamIdleTimeoutSeconds(primaryModelConfig));
ToolExecutionExecutor executor = new ToolExecutionExecutor(
toolSet, toolGuardService, approvalService, streamTracker,
toolTimeoutProperties, toolResultStorage, toolConcurrencyRegistry,
@ -943,6 +937,13 @@ public class AgentGraphBuilder {
return perSegment * (1 + vip.mate.goal.config.GoalProperties.MAX_HARD_CONTINUATIONS_CEILING) + 100;
}
static long resolveStreamIdleTimeoutSeconds(ModelConfigEntity modelConfig) {
Integer override = modelConfig != null
? modelConfig.getRequestTimeoutSeconds()
: null;
return HttpTimeouts.resolveStreamIdleTimeout(override).toSeconds();
}
CompiledGraph buildReActGraph(AgentToolSet toolSet, ChatModel chatModel, int maxIterations, String reasoningEffort) {
return buildReActGraph(toolSet, chatModel, maxIterations, reasoningEffort, null, null);
}
@ -983,16 +984,9 @@ public class AgentGraphBuilder {
contextWindowResolver.noteContextLimitError(
primaryModelConfig.getProvider(),
primaryModelConfig.getModelName(), errorMessage));
// Issue #585: drive the streaming inter-frame idle timeout from
// the per-model read-timeout knob so a stalled provider can't
// hang the body Flux after the response headers arrive. Only
// override when the model explicitly sets a value — otherwise
// the helper keeps its 180s default.
Integer perModelTimeout = primaryModelConfig.getRequestTimeoutSeconds();
if (perModelTimeout != null) {
streamingHelper.setStreamIdleTimeoutSec(perModelTimeout);
}
}
streamingHelper.setStreamIdleTimeoutSec(
resolveStreamIdleTimeoutSeconds(primaryModelConfig));
ToolExecutionExecutor executor = new ToolExecutionExecutor(
toolSet, toolGuardService, approvalService, streamTracker,
toolTimeoutProperties, toolResultStorage, toolConcurrencyRegistry,

553
mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java
View File

@ -153,6 +153,8 @@ public class ChatStreamTracker {
final List<SseEvent> buffer = new ArrayList<>();
final Object lock = new Object();
volatile boolean done;
/** Guarded by lock; once true, cleanup owns this state. */
boolean evicting;
/**
* Monotonic sequence used as the SSE protocol {@code id:} field.
* Incremented inside {@code state.lock} as each event is buffered,
@ -249,6 +251,19 @@ public class ChatStreamTracker {
}
}
/**
* Opaque lease for one exact RunState generation. Async producers should
* retain this handle so late callbacks cannot mutate a replacement run
* that happens to reuse the same conversation ID.
*/
public static final class RunHandle {
private final RunState state;
private RunHandle(RunState state) {
this.state = state;
}
}
private final ConcurrentHashMap<String, RunState> runs = new ConcurrentHashMap<>();
/**
@ -470,37 +485,43 @@ public class ChatStreamTracker {
* 注册流状态（开始生成时调用）。
* 幂等：如果已存在活跃的 RunState（Replay 与原始流共享场景），复用它而非覆盖。
*/
public void register(String conversationId) {
runs.computeIfAbsent(conversationId, RunState::new);
// 如果已存在但 done=true（上一轮残留），替换为新的
RunState state = runs.get(conversationId);
if (state != null && state.done) {
stopHeartbeat(conversationId);
RunState nextState = new RunState(conversationId);
int carried = 0;
QueuedInput queued;
while ((queued = state.messageQueue.poll()) != null) {
nextState.messageQueue.offer(queued);
carried++;
public RunHandle register(String conversationId) {
long registeredAt = System.currentTimeMillis();
RunState state = runs.compute(conversationId, (id, current) -> {
if (current == null) {
return new RunState(id);
}
runs.put(conversationId, nextState);
if (carried > 0) {
log.info("[ChatStreamTracker] Carried {} queued message(s) into next run: {}",
carried, conversationId);
synchronized (current.lock) {
if (current.evicting) {
log.info("[ChatStreamTracker] Replacing evicting run on register: {}", id);
return new RunState(id);
}
if (current.done) {
stopHeartbeat(current);
RunState nextState = new RunState(id);
int carried = 0;
QueuedInput queued;
while ((queued = current.messageQueue.poll()) != null) {
nextState.messageQueue.offer(queued);
carried++;
}
if (carried > 0) {
log.info("[ChatStreamTracker] Carried {} queued message(s) into next run: {}",
carried, id);
}
return nextState;
}
// Registration is a fresh lifecycle entrance. Refresh every
// stale-run input while holding the same lock cleanup uses to
// claim eviction, closing the former post-compute race window.
current.subscribersZeroSince = null;
current.lastEventAt = registeredAt;
if (current.stopRequested.compareAndSet(true, false)) {
log.info("[ChatStreamTracker] Reset stale stopRequested on register: {}", id);
}
}
} else if (state != null) {
// Reuse path: when complete() early-returns due to activeFluxCount > 0
// (approval replay / interrupt / any leaked flux increment), the RunState
// is kept with stopRequested still true from the previous turn. Left alone,
// the next register() would reuse it and ReasoningNode would instantly
// abort every new message with "Stop requested before LLM call".
// Reset the flag here — new registration means new user intent, and any
// still-live prior flux has already been cancelled via requestStop()'s
// disposable.dispose(), so the flag is redundant for it.
if (state.stopRequested.compareAndSet(true, false)) {
log.info("[ChatStreamTracker] Reset stale stopRequested on register: {}", conversationId);
}
}
return current;
});
// Clear the force-recycle marker on new registration — the recycle
// tombstone is meant to suppress the late doOnComplete of the
// *recycled* run only, not future turns on the same conversation. If
@ -510,8 +531,10 @@ public class ChatStreamTracker {
if (recycledConversations.remove(conversationId) != null) {
log.info("[ChatStreamTracker] Cleared recycle marker on new register: {}", conversationId);
}
startHeartbeat(conversationId);
RunHandle handle = new RunHandle(state);
startHeartbeat(state);
log.debug("Stream registered: {}", conversationId);
return handle;
}
/**
@ -524,6 +547,15 @@ public class ChatStreamTracker {
}
}
public void setDisposable(RunHandle handle, Disposable disposable) {
if (handle == null) return;
RunState state = handle.state;
synchronized (state.lock) {
if (!isCurrent(state)) return;
state.disposable = disposable;
}
}
/**
* Register an emergency-save callback for this run, invoked from {@link #onShutdown()}
* before the JVM tears down. The callback should snapshot the current accumulator
@ -536,6 +568,21 @@ public class ChatStreamTracker {
}
}
public void setEmergencySaveCallback(RunHandle handle, Runnable callback) {
if (handle == null) return;
RunState state = handle.state;
synchronized (state.lock) {
if (!isCurrent(state)) return;
state.emergencySaveCallback = callback;
}
}
private boolean isCurrent(RunState state) {
// Callers must hold state.lock so validation and mutation share one
// critical section with cleanup's evicting claim.
return !state.evicting && runs.get(state.conversationId) == state;
}
/**
* 请求停止指定会话的流。
* 取消 Flux 订阅（底层 HTTP 连接也会随之关闭），返回 true 表示确实停止了正在运行的流。
@ -603,6 +650,81 @@ public class ChatStreamTracker {
broadcast(conversationId, eventName, jsonData, false);
}
public void broadcast(RunHandle handle, String eventName, String jsonData) {
broadcast(handle, eventName, jsonData, false);
}
public void broadcast(RunHandle handle, String eventName, String jsonData, boolean skipBuffer) {
if (handle == null) return;
RunState state = handle.state;
boolean isDone = "done".equals(eventName);
boolean isAsyncTask = eventName != null && eventName.startsWith("async_task_");
boolean isHeartbeat = "heartbeat".equals(eventName);
List<SseEmitter> targets;
long eventId = 0L;
boolean forwardRelays;
synchronized (state.lock) {
if (!isCurrent(state)) return;
if (!isHeartbeat) {
state.lastEventAt = System.currentTimeMillis();
}
if (!isDone && !isAsyncTask && !isHeartbeat && state.done) {
return;
}
if ((isDone || isAsyncTask) || (!isHeartbeat && !skipBuffer)) {
eventId = ++state.nextEventId;
state.buffer.add(new SseEvent(eventId, eventName, jsonData));
if (state.buffer.size() > MAX_BUFFER_SIZE) {
trimBuffer(state.buffer);
}
}
targets = new ArrayList<>(state.subscribers);
forwardRelays = !isDone && !isAsyncTask && !isHeartbeat;
}
List<SseEmitter> dead = new ArrayList<>();
for (SseEmitter emitter : targets) {
try {
SseEmitter.SseEventBuilder event = SseEmitter.event().name(eventName).data(jsonData);
if (!isHeartbeat && !skipBuffer) {
event.id(String.valueOf(eventId));
}
emitter.send(event);
} catch (IOException | IllegalStateException e) {
dead.add(emitter);
log.debug("Removing dead subscriber for {} while sending {} event: {}",
state.conversationId, eventName, e.getMessage());
}
}
if (!dead.isEmpty()) {
synchronized (state.lock) {
if (isCurrent(state)) {
boolean removed = state.subscribers.removeAll(dead);
if (removed && state.subscribers.isEmpty()
&& !state.done && state.subscribersZeroSince == null) {
state.subscribersZeroSince = System.currentTimeMillis();
}
}
}
}
if (forwardRelays) {
List<java.util.function.BiConsumer<String, String>> relays =
eventRelays.get(state.conversationId);
if (relays != null) {
for (var relay : relays) {
try {
relay.accept(eventName, jsonData);
} catch (Exception e) {
log.debug("Event relay error for {}: {}",
state.conversationId, e.getMessage());
}
}
}
}
}
/**
* Broadcast an event to all subscribers (optionally skip buffer).
* @param skipBuffer if true, do not write to the ring buffer — used for
@ -930,6 +1052,14 @@ public class ChatStreamTracker {
return attach(conversationId, emitter, 0L);
}
public boolean attach(RunHandle handle, SseEmitter emitter) {
return attach(handle, emitter, 0L);
}
public boolean attach(RunHandle handle, SseEmitter emitter, long lastEventId) {
return handle != null && attach(handle.state, emitter, lastEventId);
}
/**
* Reconnect-aware attach: replays only events whose id &gt;
* {@code lastEventId}. Pass 0 to replay everything (fresh attach
@ -945,10 +1075,19 @@ public class ChatStreamTracker {
*/
public boolean attach(String conversationId, SseEmitter emitter, long lastEventId) {
RunState state = runs.get(conversationId);
return attach(state, emitter, lastEventId);
}
private boolean attach(RunState state, SseEmitter emitter, long lastEventId) {
if (state == null) {
return false;
}
String conversationId = state.conversationId;
synchronized (state.lock) {
if (!isCurrent(state)) {
log.info("[SSE] Attach rejected because run is being evicted: {}", conversationId);
return false;
}
// Replay buffer with id-based dedup. Each buffered event keeps its
// original (1:1) id, so the skip condition is the simple
// `id <= lastEventId`. trimBuffer no longer merges delta events,
@ -999,7 +1138,7 @@ public class ChatStreamTracker {
// Restart heartbeat so the proxy/Tomcat 60s idle timeout doesn't
// close the reconnected emitter before the async_task_* event fires.
// The scheduler self-stops once subscribers go empty (see startHeartbeat).
startHeartbeat(conversationId);
startHeartbeat(state);
return true;
}
}
@ -1042,20 +1181,37 @@ public class ChatStreamTracker {
if (state == null) {
return true;
}
return complete(state);
}
public boolean complete(RunHandle handle) {
return handle != null && complete(handle.state);
}
private boolean complete(RunState state) {
String conversationId = state.conversationId;
ScheduledFuture<?> oldHeartbeat;
synchronized (state.lock) {
if (!isCurrent(state)) {
return false;
}
state.activeFluxCount = Math.max(0, state.activeFluxCount - 1);
if (state.activeFluxCount > 0) {
log.debug("Stream partially completed (no queue drain): {} (remaining flux={})",
conversationId, state.activeFluxCount);
return false;
}
state.done = true;
oldHeartbeat = state.heartbeatFuture;
state.heartbeatFuture = null;
}
// 所有 Flux 都已完成，停止心跳，标记 done 但**不立即移除 RunState**——
// 留给 cleanupStaleRuns 在 DONE_RETENTION_MS 后异步清理。这段窗口期内
// 客户端刷新页面 attach() 能从 buffer 回放 done 事件，UI 不会卡在
// "生成中"。之前立即 runs.remove() 是 SSE 中途断开导致 done 永远丢的根源。
stopHeartbeat(conversationId);
state.done = true;
if (oldHeartbeat != null) {
oldHeartbeat.cancel(false);
}
log.debug("Stream fully completed (no queue drain): {} (kept in map for {}ms reconnect window)",
conversationId, DONE_RETENTION_MS);
return true;
@ -1075,7 +1231,11 @@ public class ChatStreamTracker {
return new CompletionResult(true, null);
}
QueuedInput consumed = null;
ScheduledFuture<?> oldHeartbeat;
synchronized (state.lock) {
if (!isCurrent(state)) {
return new CompletionResult(false, null);
}
state.activeFluxCount = Math.max(0, state.activeFluxCount - 1);
if (state.activeFluxCount > 0) {
log.debug("Stream partially completed: {} (remaining flux={}, queuePreserved={})",
@ -1084,11 +1244,15 @@ public class ChatStreamTracker {
}
// 最后一个 Flux：在同一个锁内消费排队消息（取队首）
consumed = state.messageQueue.poll();
state.done = true;
oldHeartbeat = state.heartbeatFuture;
state.heartbeatFuture = null;
}
// 锁外：停止心跳，标记 done。**不立即移除 RunState**——保留 DONE_RETENTION_MS
// 锁外：仅取消锁内摘除的旧心跳。**不立即移除 RunState**——保留 DONE_RETENTION_MS
// 让客户端可在窗口期内刷新页面通过 attach() 回放 done 事件。
stopHeartbeat(conversationId);
state.done = true;
if (oldHeartbeat != null) {
oldHeartbeat.cancel(false);
}
log.debug("Stream fully completed: {} (hasQueuedSnapshot={}, kept in map for {}ms reconnect window)",
conversationId, consumed != null, DONE_RETENTION_MS);
return new CompletionResult(true, consumed);
@ -1107,17 +1271,30 @@ public class ChatStreamTracker {
*/
public void detach(String conversationId, SseEmitter emitter) {
RunState state = runs.get(conversationId);
detach(state, emitter, false);
}
public void detach(RunHandle handle, SseEmitter emitter) {
if (handle != null) {
detach(handle.state, emitter, true);
}
}
private void detach(RunState state, SseEmitter emitter, boolean armWhenAlreadyAbsent) {
if (state == null) {
return;
}
String conversationId = state.conversationId;
synchronized (state.lock) {
state.subscribers.remove(emitter);
if (!isCurrent(state)) return;
boolean removed = state.subscribers.remove(emitter);
// When the last subscriber leaves and the run is still alive, arm
// the orphan clock — see RunState.subscribersZeroSince. The run
// is now invisible to its owner and (for webchat) unreachable, so
// cleanupStaleRuns will reclaim it after the grace window unless a
// fresh subscriber re-attaches (which clears the clock in attach()).
if (state.subscribers.isEmpty() && !state.done && state.subscribersZeroSince == null) {
if ((removed || armWhenAlreadyAbsent) && state.subscribers.isEmpty()
&& !state.done && state.subscribersZeroSince == null) {
state.subscribersZeroSince = System.currentTimeMillis();
}
}
@ -1146,45 +1323,53 @@ public class ChatStreamTracker {
*/
public void startHeartbeat(String conversationId) {
RunState state = runs.get(conversationId);
if (state == null) return;
// 避免重复启动
if (state.heartbeatFuture != null && !state.heartbeatFuture.isDone()) return;
startHeartbeat(state);
}
int intervalSec = currentHeartbeatIntervalSec(state);
state.heartbeatFuture = heartbeatScheduler.scheduleAtFixedRate(() -> {
try {
RunState s = runs.get(conversationId);
if (s == null) {
stopHeartbeat(conversationId);
return;
}
// Continue heartbeating post-done as long as someone is still listening
// (reconnected emitter waiting for late async_task_* events). Stop only
// when the run is done AND the subscribers list is empty — otherwise the
// 60s idle proxy timeout drops the reconnected emitter and async events
// never reach the client live.
if (s.done && s.subscribers.isEmpty()) {
stopHeartbeat(conversationId);
return;
}
String json;
private void startHeartbeat(RunState state) {
if (state == null) return;
String conversationId = state.conversationId;
synchronized (state.lock) {
if (!isCurrent(state)) return;
// 避免重复启动
if (state.heartbeatFuture != null && !state.heartbeatFuture.isDone()) return;
int intervalSec = currentHeartbeatIntervalSec(state);
RunHandle heartbeatHandle = new RunHandle(state);
state.heartbeatFuture = heartbeatScheduler.scheduleAtFixedRate(() -> {
try {
json = objectMapper.writeValueAsString(Map.of(
"conversationId", conversationId,
"currentPhase", safe(s.currentPhase),
"waitingReason", safe(s.waitingReason),
"runningToolName", safe(s.runningToolName),
"queueLength", s.messageQueue.size(),
"timestamp", System.currentTimeMillis()
));
boolean shouldStop;
synchronized (state.lock) {
shouldStop = !isCurrent(state)
|| (state.done && state.subscribers.isEmpty());
}
// Continue heartbeating post-done as long as someone is still listening
// (reconnected emitter waiting for late async_task_* events). Stop only
// when the run is done AND the subscribers list is empty — otherwise the
// 60s idle proxy timeout drops the reconnected emitter and async events
// never reach the client live.
if (shouldStop) {
stopHeartbeat(state);
return;
}
String json;
try {
json = objectMapper.writeValueAsString(Map.of(
"conversationId", conversationId,
"currentPhase", safe(state.currentPhase),
"waitingReason", safe(state.waitingReason),
"runningToolName", safe(state.runningToolName),
"queueLength", state.messageQueue.size(),
"timestamp", System.currentTimeMillis()
));
} catch (Exception e) {
json = "{\"conversationId\":\"" + conversationId + "\"}";
}
broadcast(heartbeatHandle, "heartbeat", json);
} catch (Exception e) {
json = "{\"conversationId\":\"" + conversationId + "\"}";
log.debug("Heartbeat error for {}: {}", conversationId, e.getMessage());
}
broadcast(conversationId, "heartbeat", json);
} catch (Exception e) {
log.debug("Heartbeat error for {}: {}", conversationId, e.getMessage());
}
}, intervalSec, intervalSec, TimeUnit.SECONDS);
}, intervalSec, intervalSec, TimeUnit.SECONDS);
}
}
/**
@ -1222,7 +1407,10 @@ public class ChatStreamTracker {
* 停止心跳定时器
*/
public void stopHeartbeat(String conversationId) {
RunState state = runs.get(conversationId);
stopHeartbeat(runs.get(conversationId));
}
private void stopHeartbeat(RunState state) {
if (state != null && state.heartbeatFuture != null) {
state.heartbeatFuture.cancel(false);
state.heartbeatFuture = null;
@ -1545,6 +1733,9 @@ public class ChatStreamTracker {
/** 已完成的 RunState 保留时间（5 分钟） */
private static final long DONE_RETENTION_MS = 5 * 60 * 1000;
/** Stale-run sweep cadence; bounds orphan eviction delay beyond the grace period. */
static final long STALE_RUN_SWEEP_INTERVAL_MS = 30_000L;
/**
* RunState 最长无活动时间。从 wall-clock {@code MAX_LIFETIME_MS=30min}
* 切换到 inactivity-based 后默认 30 min（1800s 空闲超时）：只要 agent 还在持续产事件
@ -1599,6 +1790,23 @@ public class ChatStreamTracker {
return runs.containsKey(conversationId);
}
/** Test hook — true when the current RunState owns a live heartbeat. */
boolean hasHeartbeatForTesting(String conversationId) {
RunState state = runs.get(conversationId);
if (state == null) return false;
ScheduledFuture<?> future = state.heartbeatFuture;
return future != null && !future.isCancelled();
}
/** Test hook — current generation's replay-buffer size. */
int bufferSizeForTesting(String conversationId) {
RunState state = runs.get(conversationId);
if (state == null) return 0;
synchronized (state.lock) {
return state.buffer.size();
}
}
/** Test hook — exposes the configurable timeout for assertion. */
int idleTimeoutMinutesForTesting() {
return idleTimeoutMinutes;
@ -1630,96 +1838,113 @@ public class ChatStreamTracker {
* - 订阅者清零超过 {@link #orphanGraceSeconds} 且仍在运行的孤儿 →
* 移除（webchat 无重连端点，运行对调用方不可见不可达，见 #587）
*/
@org.springframework.scheduling.annotation.Scheduled(fixedRate = 600_000)
@org.springframework.scheduling.annotation.Scheduled(
fixedDelay = STALE_RUN_SWEEP_INTERVAL_MS)
public void cleanupStaleRuns() {
long now = System.currentTimeMillis();
long idleThresholdMs = (long) idleTimeoutMinutes * 60_000L;
long orphanGraceMs = (long) orphanGraceSeconds * 1000L;
int evicted = 0;
int reclaimed = 0;
int mappingsRemoved = 0;
var iterator = runs.entrySet().iterator();
while (iterator.hasNext()) {
var entry = iterator.next();
for (var entry : runs.entrySet()) {
RunState state = entry.getValue();
long age = now - state.createdAt;
long idleMs = now - state.lastEventAt;
Long orphanSince = state.subscribersZeroSince;
long orphanMs = orphanSince != null ? now - orphanSince : -1L;
// Subscriber count under the lock so the orphan decision is
// consistent with the subscriber list (subscribersZeroSince is
// normally null whenever a subscriber is present, but a concurrent
// attach/detach could race the read — guard against that here).
int subCount;
String reason;
boolean saveBeforeEviction;
synchronized (state.lock) {
subCount = state.subscribers.size();
}
reason = null;
if (!state.evicting) {
long age = now - state.createdAt;
long idleMs = now - state.lastEventAt;
Long orphanSince = state.subscribersZeroSince;
long orphanMs = orphanSince != null ? now - orphanSince : -1L;
boolean shouldEvict = false;
String reason = null;
if (state.done && age > DONE_RETENTION_MS) {
reason = "completed and expired";
} else if (!state.done && state.subscribers.isEmpty()
&& orphanSince != null && orphanMs > orphanGraceMs) {
// Orphan: subscriber list empty longer than the grace window
// while the agent Flux is still running. Invisible + (for
// webchat) unreachable, so reclaim it instead of letting it
// burn tokens until the idle sweep (issue #587). A run that's
// actively producing events is NOT exempt — the whole point is
// nobody is watching those events.
reason = "orphaned: no subscribers for " + (orphanMs / 1000)
+ "s (grace " + orphanGraceSeconds + "s); run still active";
} else if (idleMs > idleThresholdMs) {
reason = "idle for " + (idleMs / 1000) + "s (threshold "
+ idleTimeoutMinutes + "min); total wall-clock age "
+ (age / 1000) + "s";
}
if (state.done && age > DONE_RETENTION_MS) {
shouldEvict = true;
reason = "completed and expired";
} else if (!state.done && subCount == 0 && orphanSince != null && orphanMs > orphanGraceMs) {
// Orphan: subscriber list empty longer than the grace window
// while the agent Flux is still running. Invisible + (for
// webchat) unreachable, so reclaim it instead of letting it
// burn tokens until the idle sweep (issue #587). A run that's
// actively producing events is NOT exempt — the whole point is
// nobody is watching those events.
shouldEvict = true;
reason = "orphaned: no subscribers for " + (orphanMs / 1000)
+ "s (grace " + orphanGraceSeconds + "s); run still active";
} else if (idleMs > idleThresholdMs) {
shouldEvict = true;
reason = "idle for " + (idleMs / 1000) + "s (threshold "
+ idleTimeoutMinutes + "min); total wall-clock age "
+ (age / 1000) + "s";
}
if (shouldEvict) {
// Flush any accumulated assistant content/segments BEFORE we
// dispose the run — mirrors {@link #onShutdown()} so an idle-
// timeout eviction doesn't leave the conversation with only
// the user message and no assistant trace (the round-6
// failure mode: SSE evicted mid-stream, UI refresh saw blank
// because doOnComplete never fired for the disposed Flux).
// Skip on completed runs — they already saved via the normal
// doOnComplete path.
if (!state.done) {
Runnable cb = state.emergencySaveCallback;
if (cb != null) {
try {
cb.run();
log.info("[SSE] Emergency-saved state for conversation={} before eviction",
entry.getKey());
} catch (Exception ex) {
log.warn("[SSE] Emergency save failed for conversation={}: {}",
entry.getKey(), ex.getMessage());
}
if (reason != null) {
state.evicting = true;
}
}
// 先清理资源再移除
stopHeartbeat(entry.getKey());
// Close subscriber SSE connections so an evicted run does not
// leave clients hanging in silence until their own emitter
// timeout (issue #586). Aligns the eviction path with the
// close-out sequence forceRecycle() uses.
closeSubscribers(entry.getKey());
Disposable d = state.disposable;
if (d != null && !d.isDisposed()) {
d.dispose();
saveBeforeEviction = reason != null && !state.done;
}
if (reason != null) {
boolean mappingRemoved;
try {
// Flush any accumulated assistant content/segments BEFORE we
// dispose the run — mirrors {@link #onShutdown()} so an idle-
// timeout eviction doesn't leave the conversation with only
// the user message and no assistant trace. Skip on completed
// runs — they already saved via the normal completion path.
if (saveBeforeEviction) {
Runnable cb = state.emergencySaveCallback;
if (cb != null) {
try {
cb.run();
log.info("[SSE] Emergency-saved state for conversation={} before eviction",
entry.getKey());
} catch (Exception ex) {
log.warn("[SSE] Emergency save failed for conversation={}: {}",
entry.getKey(), ex.getMessage());
}
}
}
try {
stopHeartbeat(state);
} catch (Exception ex) {
log.warn("[SSE] Heartbeat stop failed for conversation={}: {}",
entry.getKey(), ex.getMessage());
}
// Close subscriber SSE connections so an evicted run does not
// leave clients hanging until their own emitter timeout.
try {
closeSubscribers(state, false);
} catch (Exception ex) {
log.warn("[SSE] Subscriber close failed for conversation={}: {}",
entry.getKey(), ex.getMessage());
}
try {
Disposable d = state.disposable;
if (d != null && !d.isDisposed()) {
d.dispose();
}
} catch (Exception ex) {
log.warn("[SSE] Disposable teardown failed for conversation={}: {}",
entry.getKey(), ex.getMessage());
}
} finally {
mappingRemoved = runs.remove(entry.getKey(), state);
reclaimed++;
if (mappingRemoved) {
mappingsRemoved++;
}
log.warn("[SSE] Reclaimed stale RunState resources for conversation={}: {}; "
+ "mappingRemoved={}",
entry.getKey(), reason, mappingRemoved);
}
iterator.remove();
evicted++;
log.warn("[SSE] Evicted stale RunState for conversation={}: {}",
entry.getKey(), reason);
}
}
if (evicted > 0) {
log.info("[SSE] Cleanup completed: evicted {} stale RunState entries, {} remaining",
evicted, runs.size());
if (reclaimed > 0) {
log.info("[SSE] Cleanup completed: reclaimed {} stale RunState resource set(s), "
+ "removed {} map entry/entries, {} remaining",
reclaimed, mappingsRemoved, runs.size());
}
// Age out the recycled-marker map alongside RunState cleanup. Same
@ -1885,19 +2110,33 @@ public class ChatStreamTracker {
* cannot abort the loop before later subscribers are closed.
*/
public void closeSubscribers(String conversationId) {
RunState state = runs.get(conversationId);
closeSubscribers(runs.get(conversationId), true);
}
public void closeSubscribers(RunHandle handle) {
if (handle != null) {
closeSubscribers(handle.state, true);
}
}
private void closeSubscribers(RunState state, boolean requireCurrent) {
if (state == null) return;
List<SseEmitter> subscribers;
synchronized (state.lock) {
for (SseEmitter em : state.subscribers) {
try {
em.complete();
} catch (Exception ignored) {
// A subscriber that is already closed/errored must not
// prevent the rest from being closed.
}
if (requireCurrent && !isCurrent(state)) {
return;
}
subscribers = new ArrayList<>(state.subscribers);
state.subscribers.clear();
}
for (SseEmitter em : subscribers) {
try {
em.complete();
} catch (Exception ignored) {
// A subscriber that is already closed/errored must not
// prevent the rest from being closed.
}
}
}
/**

148
mateclaw-server/src/main/java/vip/mate/channel/webchat/WebChatController.java
View File

@ -46,6 +46,8 @@ import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import reactor.core.Disposable;
@ -186,6 +188,9 @@ public class WebChatController {
log.info("[WebChat] Stream: agentId={}, conversationId={}, visitor={}", agentId, conversationId, visitorId);
AtomicReference<ChatStreamTracker.RunHandle> runHandleRef = new AtomicReference<>();
AtomicBoolean disconnected = new AtomicBoolean();
// Register emitter callbacks. An SSE disconnect means this subscriber
// left, not that the agent run finished. Use detach() instead of
// complete(); complete() would prematurely mark the RunState done,
@ -193,13 +198,15 @@ public class WebChatController {
// completion when the agent Flux actually finishes.
emitter.onCompletion(() -> {
log.debug("[WebChat] SSE completed: {}", conversationId);
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
});
emitter.onTimeout(() -> {
// INFO: a timeout means the stream went idle past the SseEmitter
// budget, which is a key signal when diagnosing stream stalls.
log.info("[WebChat] SSE timeout (stream went idle past the emitter budget): {}", conversationId);
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
// Explicitly complete after timeout so the servlet container does
// not rethrow AsyncRequestTimeoutException.
emitter.complete();
@ -213,7 +220,8 @@ public class WebChatController {
} else {
log.info("[WebChat] SSE error: {} - {}", conversationId, e.getMessage());
}
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
});
sseExecutor.execute(() -> {
@ -234,13 +242,17 @@ public class WebChatController {
}
// 初始化 SSE 流跟踪
streamTracker.register(conversationId);
streamTracker.attach(conversationId, emitter);
ChatStreamTracker.RunHandle runHandle = streamTracker.register(conversationId);
runHandleRef.set(runHandle);
streamTracker.attach(runHandle, emitter);
if (disconnected.get()) {
streamTracker.detach(runHandle, emitter);
}
// Echo the effective session so the caller can persist it (especially when
// sessionId was omitted) and address the same thread on subsequent calls. The
// visitorToken must be stored by the caller and sent back on list/messages/delete.
streamTracker.broadcast(conversationId, "meta",
streamTracker.broadcast(runHandle, "meta",
"{\"sessionId\":" + escapeJson(effectiveSessionId)
+ ",\"conversationId\":" + escapeJson(conversationId)
+ ",\"visitorToken\":" + escapeJson(visitorToken) + "}");
@ -285,18 +297,19 @@ public class WebChatController {
// indicator (phase), tool execution badges (tool_start/end),
// plan-execute checklist (plan). See docs/zh/webchat.md.
if (delta.isEvent()) {
forwardVisitorEvent(conversationId, delta.eventType(), delta.eventData());
forwardVisitorEvent(runHandle, conversationId,
delta.eventType(), delta.eventData());
}
if (delta.content() != null && !delta.content().isEmpty()) {
assistantReply.append(delta.content());
if (!delta.persistenceOnly()) {
streamTracker.broadcast(conversationId, "content_delta",
streamTracker.broadcast(runHandle, "content_delta",
"{\"text\":" + escapeJson(delta.content()) + "}");
}
}
if (delta.thinking() != null && !delta.thinking().isEmpty()
&& !delta.persistenceOnly()) {
streamTracker.broadcast(conversationId, "thinking_delta",
streamTracker.broadcast(runHandle, "thinking_delta",
"{\"text\":" + escapeJson(delta.thinking()) + "}");
}
})
@ -314,7 +327,7 @@ public class WebChatController {
log.warn("[WebChat] Failed to persist assistant reply / publish event: {}",
persistErr.getMessage());
}
streamTracker.broadcast(conversationId, "done", "{\"status\":\"completed\"}");
streamTracker.broadcast(runHandle, "done", "{\"status\":\"completed\"}");
// WebChat is a pure-backend SSE channel with no re-attach
// endpoint: for third-party integrators reading until the
// server closes, `done` IS the end of the stream. Close the
@ -322,15 +335,15 @@ public class WebChatController {
// a downstream connection-pool slot for the full SseEmitter
// timeout (issue #586). The in-house web channel does NOT
// do this — it keeps emitters open for reconnect + replay.
streamTracker.closeSubscribers(conversationId);
streamTracker.complete(conversationId);
streamTracker.closeSubscribers(runHandle);
streamTracker.complete(runHandle);
})
.doOnError(e -> {
log.error("[WebChat] Stream error: {}", e.getMessage());
streamTracker.broadcast(conversationId, "error",
streamTracker.broadcast(runHandle, "error",
"{\"message\":" + escapeJson(e.getMessage()) + "}");
streamTracker.closeSubscribers(conversationId);
streamTracker.complete(conversationId);
streamTracker.closeSubscribers(runHandle);
streamTracker.complete(runHandle);
})
.subscribe();
// Bind the subscription's Disposable so requestStop() (invoked by
@ -338,12 +351,12 @@ public class WebChatController {
// the LLM stream. Without this, stopRequested is set but the underlying
// HTTP call keeps running — token burn + side-effect tools still fire.
// Mirrors ChatController#chatStream line 495.
streamTracker.setDisposable(conversationId, disposable);
streamTracker.setDisposable(runHandle, disposable);
// Wire the emergency save so an orphaned run (only subscriber
// gone) is flushed as an "interrupted" assistant message when
// the grace-period eviction reclaims it — otherwise the visitor
// would see only their own user message (issue #587).
registerEmergencySave(conversationId, assistantReply, usage, modelInfo);
registerEmergencySave(runHandle, conversationId, assistantReply, usage, modelInfo);
} catch (Exception e) {
log.error("[WebChat] Error: {}", e.getMessage(), e);
@ -1266,13 +1279,18 @@ public class WebChatController {
return emitter;
}
AtomicReference<ChatStreamTracker.RunHandle> runHandleRef = new AtomicReference<>();
AtomicBoolean disconnected = new AtomicBoolean();
emitter.onCompletion(() -> {
log.debug("[WebChat] approve SSE completed: {}", conversationId);
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
});
emitter.onTimeout(() -> {
log.info("[WebChat] approve SSE timeout (stream went idle past the emitter budget): {}", conversationId);
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
emitter.complete();
});
emitter.onError(e -> {
@ -1281,7 +1299,8 @@ public class WebChatController {
} else {
log.info("[WebChat] approve SSE error: {} - {}", conversationId, e.getMessage());
}
streamTracker.detach(conversationId, emitter);
disconnected.set(true);
streamTracker.detach(runHandleRef.get(), emitter);
});
String actor = webchatUsername(visitorId);
@ -1292,24 +1311,28 @@ public class WebChatController {
// resolveAndConsume left the already-resolved / error paths
// broadcasting into a subscriber-less tracker, so the SSE hung
// to the 10-min timeout (review #415).
streamTracker.register(conversationId);
streamTracker.attach(conversationId, emitter);
ChatStreamTracker.RunHandle runHandle = streamTracker.register(conversationId);
runHandleRef.set(runHandle);
streamTracker.attach(runHandle, emitter);
if (disconnected.get()) {
streamTracker.detach(runHandle, emitter);
}
try {
// Atomically consume the approval (DB + metadata + memory, single tx).
ResolveOutcome consumed = approvalService.resolveAndConsume(pendingId, actor);
if (consumed.consumedSnapshot() == null) {
// already resolved / not found — emit a terminal done so the
// SDK's stream listener closes cleanly instead of hanging.
broadcastApprovalResolved(conversationId, consumed);
streamTracker.broadcast(conversationId, "done",
broadcastApprovalResolved(runHandle, conversationId, consumed);
streamTracker.broadcast(runHandle, "done",
"{\"status\":\"already_resolved\"}");
streamTracker.closeSubscribers(conversationId);
streamTracker.closeSubscribers(runHandle);
return;
}
// Notify the SDK the approval flipped (clears the banner) before
// replay output starts streaming.
broadcastApprovalResolved(conversationId, consumed);
broadcastApprovalResolved(runHandle, conversationId, consumed);
PendingApproval snapshot = consumed.consumedSnapshot();
Long replayAgentId = snapshot.getAgentId() != null
@ -1317,9 +1340,9 @@ public class WebChatController {
if (replayAgentId == null) {
log.warn("[WebChat] approve: no agentId on consumed approval {}, cannot replay",
pendingId);
streamTracker.broadcast(conversationId, "done",
streamTracker.broadcast(runHandle, "done",
"{\"status\":\"error\",\"message\":\"No agent bound to approval\"}");
streamTracker.closeSubscribers(conversationId);
streamTracker.closeSubscribers(runHandle);
return;
}
@ -1339,10 +1362,10 @@ public class WebChatController {
// tool here can mislead the LLM on fallthrough).
String replayPrompt = "继续执行已批准的工具调用。";
StringBuilder assistantReply = new StringBuilder();
final int[] usage = {0, 0};
final int[] usage = {0, 0, 0, 0, 0};
final String[] modelInfo = {null, null};
streamTracker.broadcast(conversationId, "message_start",
streamTracker.broadcast(runHandle, "message_start",
"{\"role\":\"assistant\"}");
Disposable disposable = agentService.chatWithReplayStream(
@ -1362,18 +1385,19 @@ public class WebChatController {
if (provider != null) modelInfo[1] = provider.toString();
}
if (delta.isEvent()) {
forwardVisitorEvent(conversationId, delta.eventType(), delta.eventData());
forwardVisitorEvent(runHandle, conversationId,
delta.eventType(), delta.eventData());
}
if (delta.content() != null && !delta.content().isEmpty()) {
assistantReply.append(delta.content());
if (!delta.persistenceOnly()) {
streamTracker.broadcast(conversationId, "content_delta",
streamTracker.broadcast(runHandle, "content_delta",
"{\"text\":" + escapeJson(delta.content()) + "}");
}
}
if (delta.thinking() != null && !delta.thinking().isEmpty()
&& !delta.persistenceOnly()) {
streamTracker.broadcast(conversationId, "thinking_delta",
streamTracker.broadcast(runHandle, "thinking_delta",
"{\"text\":" + escapeJson(delta.thinking()) + "}");
}
})
@ -1388,31 +1412,31 @@ public class WebChatController {
} catch (Exception persistErr) {
log.warn("[WebChat] approve replay persist failed: {}", persistErr.getMessage());
}
streamTracker.broadcast(conversationId, "done",
streamTracker.broadcast(runHandle, "done",
"{\"status\":\"completed\"}");
// Close the WebChat SSE connection on the logical end of
// the replay stream — same rationale as /stream (issue #586).
streamTracker.closeSubscribers(conversationId);
streamTracker.complete(conversationId);
streamTracker.closeSubscribers(runHandle);
streamTracker.complete(runHandle);
})
.doOnError(e -> {
log.error("[WebChat] approve replay stream error: {}", e.getMessage());
streamTracker.broadcast(conversationId, "error",
streamTracker.broadcast(runHandle, "error",
"{\"message\":" + escapeJson(e.getMessage()) + "}");
streamTracker.closeSubscribers(conversationId);
streamTracker.complete(conversationId);
streamTracker.closeSubscribers(runHandle);
streamTracker.complete(runHandle);
})
.subscribe();
streamTracker.setDisposable(conversationId, disposable);
registerEmergencySave(conversationId, assistantReply, usage, modelInfo);
streamTracker.setDisposable(runHandle, disposable);
registerEmergencySave(runHandle, conversationId, assistantReply, usage, modelInfo);
} catch (Exception e) {
log.error("[WebChat] approve failed for {}: {}", conversationId, e.getMessage());
try {
streamTracker.broadcast(conversationId, "error",
streamTracker.broadcast(runHandle, "error",
"{\"message\":" + escapeJson(e.getMessage()) + "}");
} catch (Exception ignored) {}
streamTracker.closeSubscribers(conversationId);
streamTracker.complete(conversationId);
streamTracker.closeSubscribers(runHandle);
streamTracker.complete(runHandle);
}
});
audit(channel, visitorId, "webchat.approve-approval", conversationId,
@ -1454,16 +1478,32 @@ public class WebChatController {
private void broadcastApprovalResolved(String conversationId, ResolveOutcome outcome) {
try {
streamTracker.broadcast(conversationId, "tool_approval_resolved",
objectMapper.writeValueAsString(Map.of(
"pendingId", outcome.pendingId(),
"decision", outcome.decision() != null ? outcome.decision() : "",
"toolName", outcome.toolName() != null ? outcome.toolName() : "")));
approvalResolvedJson(outcome));
} catch (Exception e) {
log.debug("[WebChat] approval_resolved broadcast failed for {}: {}",
outcome.pendingId(), e.getMessage());
}
}
private void broadcastApprovalResolved(ChatStreamTracker.RunHandle runHandle,
String conversationId,
ResolveOutcome outcome) {
try {
streamTracker.broadcast(runHandle, "tool_approval_resolved",
approvalResolvedJson(outcome));
} catch (Exception e) {
log.debug("[WebChat] approval_resolved broadcast failed for {} in {}: {}",
outcome.pendingId(), conversationId, e.getMessage());
}
}
private String approvalResolvedJson(ResolveOutcome outcome) throws IOException {
return objectMapper.writeValueAsString(Map.of(
"pendingId", outcome.pendingId(),
"decision", outcome.decision() != null ? outcome.decision() : "",
"toolName", outcome.toolName() != null ? outcome.toolName() : ""));
}
/**
* Regenerate the last assistant reply.
* <p>
@ -1701,9 +1741,10 @@ public class WebChatController {
* @param usage [prompt, completion, cacheRead, cacheWrite, reasoning]
* @param modelInfo [runtimeModel, runtimeProvider]
*/
private void registerEmergencySave(String conversationId, StringBuilder assistantReply,
private void registerEmergencySave(ChatStreamTracker.RunHandle runHandle,
String conversationId, StringBuilder assistantReply,
int[] usage, String[] modelInfo) {
streamTracker.setEmergencySaveCallback(conversationId, () -> {
streamTracker.setEmergencySaveCallback(runHandle, () -> {
try {
String reply = assistantReply.toString();
if (reply.isBlank()) {
@ -1968,7 +2009,10 @@ public class WebChatController {
* <p>Backward compat: visitors / SDKs that don't know these event types
* silently ignore them per the SSE spec.
*/
private void forwardVisitorEvent(String conversationId, String eventType, Map<String, Object> data) {
private void forwardVisitorEvent(ChatStreamTracker.RunHandle runHandle,
String conversationId,
String eventType,
Map<String, Object> data) {
if (eventType == null || data == null) return;
Map<String, Object> payload;
String sseName;
@ -2011,7 +2055,7 @@ public class WebChatController {
}
try {
String json = objectMapper.writeValueAsString(payload);
streamTracker.broadcast(conversationId, sseName, json);
streamTracker.broadcast(runHandle, sseName, json);
} catch (Exception e) {
log.debug("[WebChat] Failed to serialize visitor event {} for {}: {}",
eventType, conversationId, e.getMessage());

24
mateclaw-server/src/test/java/vip/mate/agent/AgentGraphBuilderStreamIdleTimeoutTest.java Normal file
View File

@ -0,0 +1,24 @@
package vip.mate.agent;
import org.junit.jupiter.api.Test;
import vip.mate.llm.model.ModelConfigEntity;
import static org.junit.jupiter.api.Assertions.assertEquals;
class AgentGraphBuilderStreamIdleTimeoutTest {
@Test
void normalizesStreamIdleTimeoutOverrides() {
assertEquals(180, AgentGraphBuilder.resolveStreamIdleTimeoutSeconds(null));
assertEquals(180, AgentGraphBuilder.resolveStreamIdleTimeoutSeconds(modelConfig(null)));
assertEquals(180, AgentGraphBuilder.resolveStreamIdleTimeoutSeconds(modelConfig(0)));
assertEquals(180, AgentGraphBuilder.resolveStreamIdleTimeoutSeconds(modelConfig(-30)));
assertEquals(600, AgentGraphBuilder.resolveStreamIdleTimeoutSeconds(modelConfig(600)));
}
private static ModelConfigEntity modelConfig(Integer requestTimeoutSeconds) {
ModelConfigEntity modelConfig = new ModelConfigEntity();
modelConfig.setRequestTimeoutSeconds(requestTimeoutSeconds);
return modelConfig;
}
}

33
mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerCloseSubscribersTest.java
View File

@ -5,6 +5,10 @@ import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
@ -134,4 +138,33 @@ class ChatStreamTrackerCloseSubscribersTest {
"the live subscriber must still be closed even though a dead " +
"subscriber threw on complete()");
}
@Test
@DisplayName("Emitter completion runs outside the RunState lock")
void closeSubscribersCompletesOutsideStateLock() {
ChatStreamTracker tracker = newTracker();
String cid = "close-outside-lock";
ChatStreamTracker.RunHandle handle = tracker.register(cid);
AtomicBoolean concurrentAttachSucceeded = new AtomicBoolean();
SseEmitter emitter = new SseEmitter() {
@Override
public void complete() {
CompletableFuture<Boolean> attach = CompletableFuture.supplyAsync(
() -> tracker.attach(handle, new SseEmitter()));
try {
concurrentAttachSucceeded.set(attach.get(1, TimeUnit.SECONDS));
} catch (Exception ignored) {
concurrentAttachSucceeded.set(false);
}
super.complete();
}
};
assertTrue(tracker.attach(handle, emitter));
tracker.closeSubscribers(handle);
assertTrue(concurrentAttachSucceeded.get(),
"complete() must not run while closeSubscribers owns the state lock");
}
}

90
mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerDetachSemanticsTest.java
View File

@ -5,6 +5,14 @@ import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.assertEquals;
@ -30,6 +38,38 @@ class ChatStreamTrackerDetachSemanticsTest {
return new ChatStreamTracker(new ObjectMapper());
}
@SuppressWarnings("unchecked")
private void pauseHeartbeatCancellation(ChatStreamTracker tracker,
String conversationId,
CountDownLatch cancellationStarted,
CountDownLatch releaseCancellation) throws Exception {
Field runsField = ChatStreamTracker.class.getDeclaredField("runs");
runsField.setAccessible(true);
Map<String, ChatStreamTracker.RunState> runs =
(Map<String, ChatStreamTracker.RunState>) runsField.get(tracker);
ChatStreamTracker.RunState state = runs.get(conversationId);
ScheduledFuture<?> original = state.heartbeatFuture;
ScheduledFuture<?> blocking = (ScheduledFuture<?>) Proxy.newProxyInstance(
ScheduledFuture.class.getClassLoader(),
new Class<?>[]{ScheduledFuture.class},
(proxy, method, args) -> {
if ("cancel".equals(method.getName())) {
cancellationStarted.countDown();
if (!releaseCancellation.await(2, TimeUnit.SECONDS)) {
throw new AssertionError("heartbeat cancellation release timed out");
}
}
try {
return method.invoke(original, args);
} catch (InvocationTargetException e) {
throw e.getCause();
}
});
synchronized (state.lock) {
state.heartbeatFuture = blocking;
}
}
@Test
@DisplayName("detach() leaves the run running — isRunning() stays true")
void detachKeepsRunRunning() {
@ -108,4 +148,54 @@ class ChatStreamTrackerDetachSemanticsTest {
assertTrue(tracker.isRunning("present"));
assertEquals(0, tracker.getAllSnapshot().getFirst().subscriberCount());
}
@Test
@DisplayName("complete preserves a heartbeat started by a post-done attach")
void completeDoesNotCancelPostDoneAttachHeartbeat() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "complete-heartbeat-handoff";
tracker.register(cid);
tracker.incrementFlux(cid);
CountDownLatch cancellationStarted = new CountDownLatch(1);
CountDownLatch releaseCancellation = new CountDownLatch(1);
pauseHeartbeatCancellation(tracker, cid, cancellationStarted, releaseCancellation);
CompletableFuture<Boolean> completion =
CompletableFuture.supplyAsync(() -> tracker.complete(cid));
try {
assertTrue(cancellationStarted.await(1, TimeUnit.SECONDS));
assertTrue(tracker.attach(cid, new SseEmitter()));
} finally {
releaseCancellation.countDown();
}
assertTrue(completion.get(2, TimeUnit.SECONDS));
assertTrue(tracker.hasHeartbeatForTesting(cid),
"completion must cancel only the heartbeat detached before post-done attach");
}
@Test
@DisplayName("queue-draining completion preserves a heartbeat started by a post-done attach")
void queueDrainDoesNotCancelPostDoneAttachHeartbeat() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "queue-drain-heartbeat-handoff";
tracker.register(cid);
tracker.incrementFlux(cid);
CountDownLatch cancellationStarted = new CountDownLatch(1);
CountDownLatch releaseCancellation = new CountDownLatch(1);
pauseHeartbeatCancellation(tracker, cid, cancellationStarted, releaseCancellation);
CompletableFuture<ChatStreamTracker.CompletionResult> completion =
CompletableFuture.supplyAsync(() -> tracker.completeAndConsumeIfLast(cid));
try {
assertTrue(cancellationStarted.await(1, TimeUnit.SECONDS));
assertTrue(tracker.attach(cid, new SseEmitter()));
} finally {
releaseCancellation.countDown();
}
assertTrue(completion.get(2, TimeUnit.SECONDS).allDone());
assertTrue(tracker.hasHeartbeatForTesting(cid),
"queue drain must cancel only the heartbeat detached before post-done attach");
}
}

305
mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerOrphanPolicyTest.java
View File

@ -4,11 +4,21 @@ import com.fasterxml.jackson.databind.ObjectMapper;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.Disposable;
import java.io.IOException;
import java.lang.reflect.Field;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertNotNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
/**
@ -23,6 +33,32 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
*/
class ChatStreamTrackerOrphanPolicyTest {
private static final class RecordingDisposable implements Disposable {
private final AtomicBoolean disposed = new AtomicBoolean();
private final boolean throwOnDispose;
private RecordingDisposable() {
this(false);
}
private RecordingDisposable(boolean throwOnDispose) {
this.throwOnDispose = throwOnDispose;
}
@Override
public void dispose() {
disposed.set(true);
if (throwOnDispose) {
throw new IllegalStateException("dispose failed");
}
}
@Override
public boolean isDisposed() {
return disposed.get();
}
}
private ChatStreamTracker newTracker() {
ChatStreamTracker t = new ChatStreamTracker(new ObjectMapper());
t.setIdleTimeoutMinutesForTesting(30); // keep the idle bucket out of the way
@ -30,6 +66,27 @@ class ChatStreamTrackerOrphanPolicyTest {
return t;
}
private static CompletableFuture<Void> startPausedCleanup(
ChatStreamTracker tracker,
String conversationId,
CountDownLatch claimed,
CountDownLatch release) {
tracker.setEmergencySaveCallback(conversationId, () -> {
claimed.countDown();
try {
if (!release.await(2, TimeUnit.SECONDS)) {
throw new AssertionError("cleanup release latch timed out");
}
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new AssertionError(e);
}
});
tracker.backdateOrphanForTesting(
conversationId, System.currentTimeMillis() - 3_000L);
return CompletableFuture.runAsync(tracker::cleanupStaleRuns);
}
@Test
@DisplayName("Orphan run (no subscribers past grace) is evicted")
void orphanRunEvictedAfterGrace() {
@ -146,4 +203,252 @@ class ChatStreamTrackerOrphanPolicyTest {
assertTrue(tracker.hasRunStateForTesting(cid),
"re-attach clears the orphan clock even if it was backdated");
}
@Test
@DisplayName("Attach is rejected after cleanup atomically claims an orphan")
void attachRejectedAfterEvictionClaim() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "orphan-claim";
tracker.register(cid);
tracker.incrementFlux(cid);
CountDownLatch claimed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
CompletableFuture<Void> cleanup = startPausedCleanup(tracker, cid, claimed, release);
assertTrue(claimed.await(1, TimeUnit.SECONDS));
try {
assertFalse(tracker.attach(cid, new SseEmitter()),
"attach must not report success after eviction is claimed");
} finally {
release.countDown();
cleanup.get(2, TimeUnit.SECONDS);
}
assertFalse(tracker.hasRunStateForTesting(cid));
}
@Test
@DisplayName("Old cleanup cannot remove or close a replacement run")
void claimedCleanupDoesNotTouchReplacementRun() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "orphan-replacement";
tracker.register(cid);
tracker.incrementFlux(cid);
CountDownLatch claimed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
CompletableFuture<Void> cleanup = startPausedCleanup(tracker, cid, claimed, release);
assertTrue(claimed.await(1, TimeUnit.SECONDS));
SseEmitter replacementEmitter = new SseEmitter();
try {
tracker.register(cid);
tracker.incrementFlux(cid);
assertTrue(tracker.attach(cid, replacementEmitter));
assertTrue(tracker.hasHeartbeatForTesting(cid));
} finally {
release.countDown();
cleanup.get(2, TimeUnit.SECONDS);
}
assertTrue(tracker.hasRunStateForTesting(cid));
assertTrue(tracker.isRunning(cid));
assertTrue(tracker.hasHeartbeatForTesting(cid));
assertDoesNotThrow(() -> replacementEmitter.send(
SseEmitter.event().name("probe").data("still-open")));
}
@Test
@DisplayName("A stale emitter callback cannot orphan a replacement run")
void staleDetachDoesNotArmReplacementOrphanClock() {
ChatStreamTracker tracker = newTracker();
String cid = "stale-detach";
tracker.register(cid);
tracker.incrementFlux(cid);
SseEmitter oldEmitter = new SseEmitter();
assertTrue(tracker.attach(cid, oldEmitter));
tracker.complete(cid);
tracker.register(cid);
tracker.incrementFlux(cid);
// A delayed onCompletion/onTimeout callback belongs to the prior
// generation. It must not arm the fresh state's orphan clock merely
// because that new state has not attached its own emitter yet.
tracker.detach(cid, oldEmitter);
tracker.setOrphanGraceSecondsForTesting(-1);
tracker.cleanupStaleRuns();
assertTrue(tracker.hasRunStateForTesting(cid));
assertTrue(tracker.isRunning(cid));
}
@Test
@DisplayName("Register atomically refreshes a reused run before cleanup can claim it")
void registerRefreshesReusedRunLifecycle() {
ChatStreamTracker tracker = newTracker();
String cid = "register-refresh";
tracker.register(cid);
tracker.incrementFlux(cid);
tracker.backdateOrphanForTesting(cid, System.currentTimeMillis() - 3_000L);
ChatStreamTracker.RunHandle handle = tracker.register(cid);
tracker.cleanupStaleRuns();
assertTrue(tracker.hasRunStateForTesting(cid));
assertTrue(tracker.attach(handle, new SseEmitter()));
}
@Test
@DisplayName("Late callbacks from an old handle cannot mutate a replacement run")
void oldHandleCannotMutateReplacementRun() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "old-handle";
ChatStreamTracker.RunHandle oldHandle = tracker.register(cid);
tracker.incrementFlux(cid);
CountDownLatch claimed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
CompletableFuture<Void> cleanup = startPausedCleanup(tracker, cid, claimed, release);
assertTrue(claimed.await(1, TimeUnit.SECONDS));
ChatStreamTracker.RunHandle replacementHandle = tracker.register(cid);
SseEmitter replacementEmitter = new SseEmitter();
try {
assertTrue(tracker.attach(replacementHandle, replacementEmitter));
} finally {
release.countDown();
cleanup.get(2, TimeUnit.SECONDS);
}
int bufferBefore = tracker.bufferSizeForTesting(cid);
assertFalse(tracker.attach(oldHandle, new SseEmitter()));
tracker.broadcast(oldHandle, "content_delta", "{\"text\":\"stale\"}");
tracker.closeSubscribers(oldHandle);
tracker.complete(oldHandle);
assertEquals(bufferBefore, tracker.bufferSizeForTesting(cid));
assertTrue(tracker.isRunning(cid));
assertTrue(tracker.hasHeartbeatForTesting(cid));
assertDoesNotThrow(() -> replacementEmitter.send(
SseEmitter.event().name("probe").data("still-open")));
}
@Test
@DisplayName("Cleanup and stale handles cannot touch a replacement disposable")
void cleanupDisposesOnlyClaimedStateDisposable() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "disposable-generation";
ChatStreamTracker.RunHandle oldHandle = tracker.register(cid);
tracker.incrementFlux(cid);
RecordingDisposable oldDisposable = new RecordingDisposable();
tracker.setDisposable(oldHandle, oldDisposable);
CountDownLatch claimed = new CountDownLatch(1);
CountDownLatch release = new CountDownLatch(1);
CompletableFuture<Void> cleanup = startPausedCleanup(tracker, cid, claimed, release);
assertTrue(claimed.await(1, TimeUnit.SECONDS));
ChatStreamTracker.RunHandle replacementHandle = tracker.register(cid);
RecordingDisposable replacementDisposable = new RecordingDisposable();
RecordingDisposable staleLateDisposable = new RecordingDisposable();
AtomicInteger replacementSaveCount = new AtomicInteger();
AtomicInteger staleSaveCount = new AtomicInteger();
tracker.setDisposable(replacementHandle, replacementDisposable);
tracker.setDisposable(oldHandle, staleLateDisposable);
tracker.setEmergencySaveCallback(replacementHandle, replacementSaveCount::incrementAndGet);
tracker.setEmergencySaveCallback(oldHandle, staleSaveCount::incrementAndGet);
release.countDown();
cleanup.get(2, TimeUnit.SECONDS);
assertTrue(oldDisposable.isDisposed());
assertFalse(replacementDisposable.isDisposed());
assertFalse(staleLateDisposable.isDisposed());
assertTrue(tracker.hasRunStateForTesting(cid));
tracker.backdateOrphanForTesting(cid, System.currentTimeMillis() - 3_000L);
tracker.cleanupStaleRuns();
assertEquals(1, replacementSaveCount.get());
assertEquals(0, staleSaveCount.get());
assertTrue(replacementDisposable.isDisposed());
assertFalse(staleLateDisposable.isDisposed());
}
@Test
@DisplayName("A throwing disposable cannot leave an evicting tombstone mapped")
void throwingDisposableStillRemovesClaimedState() {
ChatStreamTracker tracker = newTracker();
String cid = "throwing-disposable";
ChatStreamTracker.RunHandle handle = tracker.register(cid);
tracker.incrementFlux(cid);
tracker.setDisposable(handle, new RecordingDisposable(true));
tracker.backdateOrphanForTesting(cid, System.currentTimeMillis() - 3_000L);
assertDoesNotThrow(tracker::cleanupStaleRuns);
assertFalse(tracker.hasRunStateForTesting(cid));
}
@Test
@DisplayName("Broadcast send failure arms orphan cleanup when the last subscriber is removed")
@SuppressWarnings("unchecked")
void sendFailureArmsOrphanCleanup() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "send-failure-orphan";
ChatStreamTracker.RunHandle handle = tracker.register(cid);
tracker.incrementFlux(cid);
SseEmitter failingEmitter = new SseEmitter() {
@Override
public void send(SseEventBuilder builder) throws IOException {
throw new IOException("client disconnected");
}
};
assertTrue(tracker.attach(handle, failingEmitter));
tracker.broadcast(handle, "content_delta", "{\"text\":\"still-running\"}");
Field runsField = ChatStreamTracker.class.getDeclaredField("runs");
runsField.setAccessible(true);
Map<String, ChatStreamTracker.RunState> runs =
(Map<String, ChatStreamTracker.RunState>) runsField.get(tracker);
ChatStreamTracker.RunState state = runs.get(cid);
synchronized (state.lock) {
assertNotNull(state.subscribersZeroSince,
"removing the final dead subscriber must arm the orphan clock");
state.subscribersZeroSince = System.currentTimeMillis() - 3_000L;
}
tracker.cleanupStaleRuns();
assertFalse(tracker.hasRunStateForTesting(cid));
}
@Test
@DisplayName("Exact-handle detach arms orphan cleanup before an emitter was attached")
@SuppressWarnings("unchecked")
void exactDetachArmsOrphanWithoutRemovingEmitter() throws Exception {
ChatStreamTracker tracker = newTracker();
String cid = "pre-attach-exact-detach";
ChatStreamTracker.RunHandle handle = tracker.register(cid);
tracker.incrementFlux(cid);
tracker.detach(handle, new SseEmitter());
Field runsField = ChatStreamTracker.class.getDeclaredField("runs");
runsField.setAccessible(true);
Map<String, ChatStreamTracker.RunState> runs =
(Map<String, ChatStreamTracker.RunState>) runsField.get(tracker);
ChatStreamTracker.RunState state = runs.get(cid);
synchronized (state.lock) {
assertNotNull(state.subscribersZeroSince,
"an exact disconnect must arm even when attach never added the emitter");
state.subscribersZeroSince = System.currentTimeMillis() - 3_000L;
}
tracker.cleanupStaleRuns();
assertFalse(tracker.hasRunStateForTesting(cid));
}
}

94
mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatApprovalInteractionTest.java
View File

@ -3,11 +3,17 @@ package vip.mate.channel.webchat;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.test.context.TestPropertySource;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import vip.mate.MateClawApplication;
import vip.mate.agent.AgentService;
import vip.mate.agent.model.AgentEntity;
import vip.mate.approval.ApprovalWorkflowService;
import vip.mate.approval.PendingApproval;
import vip.mate.channel.web.ChatStreamTracker;
@ -15,8 +21,12 @@ import vip.mate.channel.webchat.WebChatController.WebChatCreateSessionRequest;
import vip.mate.common.result.R;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
/**
* Verifies ISSUE #413 P1: the WebChat (API-Key) channel can now resolve tool
@ -36,7 +46,8 @@ import static org.assertj.core.api.Assertions.assertThat;
"spring.datasource.url=jdbc:h2:mem:webchat_approve_${random.uuid};MODE=MySQL;DATABASE_TO_LOWER=TRUE;CASE_INSENSITIVE_IDENTIFIERS=TRUE;DB_CLOSE_DELAY=-1",
"spring.ai.dashscope.api-key=test-key",
"spring.main.web-application-type=none",
"mateclaw.jwt.secret=webchat-it-secret-0123456789"
"mateclaw.jwt.secret=webchat-it-secret-0123456789",
"mateclaw.webchat.orphan-grace-sec=-1"
})
class WebChatApprovalInteractionTest {
@ -49,6 +60,7 @@ class WebChatApprovalInteractionTest {
@Autowired private ApprovalWorkflowService approvalService;
@Autowired private ChatStreamTracker streamTracker;
@Autowired private JdbcTemplate jdbc;
@MockBean private AgentService agentService;
@BeforeEach
void setUp() {
@ -64,6 +76,10 @@ class WebChatApprovalInteractionTest {
"workspace_id, create_time, update_time, deleted) " +
"VALUES (?, 'wc', 'webchat', ?, ?, TRUE, 1, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, 0)",
CHANNEL_ID, AGENT_ID, "{\"api_key\":\"" + API_KEY + "\"}");
AgentEntity agent = new AgentEntity();
agent.setId(AGENT_ID);
agent.setWorkspaceId(1L);
Mockito.when(agentService.getAgent(AGENT_ID)).thenReturn(agent);
}
private WebChatCreateSessionRequest req(String visitorId, String sessionId) {
@ -220,4 +236,80 @@ class WebChatApprovalInteractionTest {
assertThat(r.getCode()).isEqualTo(200);
assertThat(r.getData().get("stopped")).isEqualTo(Boolean.FALSE);
}
@Test
@DisplayName("disconnect before approval worker registration arms orphan cleanup")
void disconnectBeforeApprovalWorkerRegistrationIsNotLost() throws Exception {
String visitorId = "visitor-pre-register-disconnect";
String sessionId = "s1";
String pendingId = seedPending(visitorId, sessionId);
String cid = WebChatController.deriveConversationId(API_KEY, visitorId, sessionId);
Mockito.when(agentService.chatWithReplayStream(
eq(AGENT_ID), anyString(), eq(cid), anyString(), anyString(), any()))
.thenReturn(Flux.never());
WebChatDisconnectTestSupport.QueuedExecutorService queued =
new WebChatDisconnectTestSupport.QueuedExecutorService();
ExecutorService original = WebChatDisconnectTestSupport.swapExecutor(controller, queued);
try {
SseEmitter emitter = controller.approveSession(
API_KEY, tokenFor(visitorId), visitorId, sessionId, pendingId);
WebChatDisconnectTestSupport.fireCompletion(emitter);
queued.runNext();
} finally {
WebChatDisconnectTestSupport.swapExecutor(controller, original);
}
streamTracker.cleanupStaleRuns();
assertThat(streamTracker.streamExistsOnThisNode(cid))
.as("an approval disconnect observed before registration must arm orphan cleanup")
.isFalse();
}
@Test
@DisplayName("approval replay persists full usage metadata from _usage_final")
void approvalReplayPersistsFullUsageMetadata() throws Exception {
String visitorId = "visitor-usage-final";
String sessionId = "s1";
String pendingId = seedPending(visitorId, sessionId);
String cid = WebChatController.deriveConversationId(API_KEY, visitorId, sessionId);
Mockito.when(agentService.chatWithReplayStream(
eq(AGENT_ID), anyString(), eq(cid), anyString(), anyString(), any()))
.thenReturn(Flux.just(
AgentService.StreamDelta.event("_usage_final", Map.of(
"promptTokens", 11,
"completionTokens", 7,
"cacheReadTokens", 3,
"cacheWriteTokens", 2,
"reasoningTokens", 5,
"runtimeModelName", "mock-replay-model",
"runtimeProviderId", "mock-provider")),
new AgentService.StreamDelta("approved reply", null)));
WebChatDisconnectTestSupport.QueuedExecutorService queued =
new WebChatDisconnectTestSupport.QueuedExecutorService();
ExecutorService original = WebChatDisconnectTestSupport.swapExecutor(controller, queued);
try {
controller.approveSession(API_KEY, tokenFor(visitorId), visitorId, sessionId, pendingId);
queued.runNext();
} finally {
WebChatDisconnectTestSupport.swapExecutor(controller, original);
}
Map<String, Object> row = jdbc.queryForMap(
"SELECT content, prompt_tokens, completion_tokens, cache_read_tokens, " +
"cache_write_tokens, reasoning_tokens, runtime_model, runtime_provider " +
"FROM mate_message WHERE conversation_id = ? AND role = 'assistant' " +
"ORDER BY create_time DESC LIMIT 1",
cid);
assertThat(row.get("CONTENT")).isEqualTo("approved reply");
assertThat(((Number) row.get("PROMPT_TOKENS")).intValue()).isEqualTo(11);
assertThat(((Number) row.get("COMPLETION_TOKENS")).intValue()).isEqualTo(7);
assertThat(((Number) row.get("CACHE_READ_TOKENS")).intValue()).isEqualTo(3);
assertThat(((Number) row.get("CACHE_WRITE_TOKENS")).intValue()).isEqualTo(2);
assertThat(((Number) row.get("REASONING_TOKENS")).intValue()).isEqualTo(5);
assertThat(row.get("RUNTIME_MODEL")).isEqualTo("mock-replay-model");
assertThat(row.get("RUNTIME_PROVIDER")).isEqualTo("mock-provider");
}
}

79
mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatDisconnectTestSupport.java Normal file
View File

@ -0,0 +1,79 @@
package vip.mate.channel.webchat;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.lang.reflect.Field;
import java.util.ArrayDeque;
import java.util.List;
import java.util.Queue;
import java.util.concurrent.AbstractExecutorService;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;
final class WebChatDisconnectTestSupport {
private WebChatDisconnectTestSupport() {
}
static ExecutorService swapExecutor(WebChatController controller,
ExecutorService replacement) throws Exception {
Field field = WebChatController.class.getDeclaredField("sseExecutor");
field.setAccessible(true);
ExecutorService original = (ExecutorService) field.get(controller);
field.set(controller, replacement);
return original;
}
static void fireCompletion(SseEmitter emitter) throws Exception {
Field field = ResponseBodyEmitter.class.getDeclaredField("completionCallback");
field.setAccessible(true);
((Runnable) field.get(emitter)).run();
}
static final class QueuedExecutorService extends AbstractExecutorService {
private final Queue<Runnable> tasks = new ArrayDeque<>();
private boolean shutdown;
@Override
public void shutdown() {
shutdown = true;
}
@Override
public List<Runnable> shutdownNow() {
shutdown = true;
var remaining = List.copyOf(tasks);
tasks.clear();
return remaining;
}
@Override
public boolean isShutdown() {
return shutdown;
}
@Override
public boolean isTerminated() {
return shutdown && tasks.isEmpty();
}
@Override
public boolean awaitTermination(long timeout, TimeUnit unit) {
return isTerminated();
}
@Override
public void execute(Runnable command) {
tasks.add(command);
}
void runNext() {
Runnable task = tasks.poll();
if (task == null) {
throw new AssertionError("expected a queued SSE worker");
}
task.run();
}
}
}

39
mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatStreamE2ETest.java
View File

@ -9,10 +9,13 @@ import org.springframework.boot.test.mock.mockito.MockBean;
import org.springframework.boot.test.web.server.LocalServerPort;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.test.context.TestPropertySource;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import vip.mate.MateClawApplication;
import vip.mate.agent.AgentService;
import vip.mate.agent.model.AgentEntity;
import vip.mate.channel.web.ChatStreamTracker;
import vip.mate.channel.webchat.WebChatController.WebChatRequest;
import java.io.BufferedReader;
import java.io.IOException;
@ -28,6 +31,7 @@ import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ExecutorService;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.ArgumentMatchers.any;
@ -75,6 +79,7 @@ import static org.mockito.ArgumentMatchers.isNull;
"spring.datasource.url=jdbc:h2:mem:webchat_stream_e2e_${random.uuid};MODE=MySQL;DATABASE_TO_LOWER=TRUE;CASE_INSENSITIVE_IDENTIFIERS=TRUE;DB_CLOSE_DELAY=-1",
"spring.ai.dashscope.api-key=test-key",
"mateclaw.jwt.secret=webchat-it-secret-0123456789",
"mateclaw.webchat.orphan-grace-sec=-1",
"mateclaw.feature-flag.refresh-ms=999999"
})
class WebChatStreamE2ETest {
@ -87,6 +92,8 @@ class WebChatStreamE2ETest {
@LocalServerPort private int port;
@Autowired private JdbcTemplate jdbc;
@Autowired private WebChatController controller;
@Autowired private ChatStreamTracker streamTracker;
/** Replaced with a Mockito mock; tests stub the two methods /stream calls. */
@MockBean private AgentService agentService;
@ -248,6 +255,38 @@ class WebChatStreamE2ETest {
assertThat(lastAssistantContent(cid)).isEqualTo("Hello world!");
}
@Test
@DisplayName("disconnect before chat worker registration arms orphan cleanup")
void disconnectBeforeChatWorkerRegistrationIsNotLost() throws Exception {
org.mockito.Mockito.when(agentService.chatStructuredStream(
eq(AGENT_ID), anyString(), anyString(), anyString(), isNull(), any()))
.thenReturn(Flux.never());
String visitorId = "vE2E-pre-register-disconnect";
String sessionId = "pre-register";
String cid = WebChatController.deriveConversationId(API_KEY, visitorId, sessionId);
WebChatRequest request = new WebChatRequest();
request.setMessage("keep running");
request.setVisitorId(visitorId);
request.setSessionId(sessionId);
WebChatDisconnectTestSupport.QueuedExecutorService queued =
new WebChatDisconnectTestSupport.QueuedExecutorService();
ExecutorService original = WebChatDisconnectTestSupport.swapExecutor(controller, queued);
try {
SseEmitter emitter = controller.chatStream(API_KEY, request);
WebChatDisconnectTestSupport.fireCompletion(emitter);
queued.runNext();
} finally {
WebChatDisconnectTestSupport.swapExecutor(controller, original);
}
streamTracker.cleanupStaleRuns();
assertThat(streamTracker.streamExistsOnThisNode(cid))
.as("a disconnect observed before registration must arm exact-run orphan cleanup")
.isFalse();
}
@Test
@DisplayName("multi-chunk reply: thinking + content + usage event all broadcast; persisted content is content-only")
void multiChunkReply() throws Exception {