diff --git a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java index 0e0c7edf..001c5ca4 100644 --- a/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java +++ b/mateclaw-server/src/main/java/vip/mate/agent/AgentGraphBuilder.java @@ -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, diff --git a/mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java b/mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java index 94edb68c..43a31b83 100644 --- a/mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java +++ b/mateclaw-server/src/main/java/vip/mate/channel/web/ChatStreamTracker.java @@ -153,6 +153,8 @@ public class ChatStreamTracker { final List 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 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 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 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> 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 > * {@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 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. + } + } } /** diff --git a/mateclaw-server/src/main/java/vip/mate/channel/webchat/WebChatController.java b/mateclaw-server/src/main/java/vip/mate/channel/webchat/WebChatController.java index a6965498..f9a7921f 100644 --- a/mateclaw-server/src/main/java/vip/mate/channel/webchat/WebChatController.java +++ b/mateclaw-server/src/main/java/vip/mate/channel/webchat/WebChatController.java @@ -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 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 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. *

@@ -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 { *

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 data) { + private void forwardVisitorEvent(ChatStreamTracker.RunHandle runHandle, + String conversationId, + String eventType, + Map data) { if (eventType == null || data == null) return; Map 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()); diff --git a/mateclaw-server/src/test/java/vip/mate/agent/AgentGraphBuilderStreamIdleTimeoutTest.java b/mateclaw-server/src/test/java/vip/mate/agent/AgentGraphBuilderStreamIdleTimeoutTest.java new file mode 100644 index 00000000..d0830aa4 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/agent/AgentGraphBuilderStreamIdleTimeoutTest.java @@ -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; + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerCloseSubscribersTest.java b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerCloseSubscribersTest.java index afd2c88a..13d8ecd2 100644 --- a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerCloseSubscribersTest.java +++ b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerCloseSubscribersTest.java @@ -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 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"); + } } diff --git a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerDetachSemanticsTest.java b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerDetachSemanticsTest.java index 52f0ef95..56a7e71c 100644 --- a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerDetachSemanticsTest.java +++ b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerDetachSemanticsTest.java @@ -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 runs = + (Map) 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 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 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"); + } } diff --git a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerOrphanPolicyTest.java b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerOrphanPolicyTest.java index dd4f78a0..56d47f4e 100644 --- a/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerOrphanPolicyTest.java +++ b/mateclaw-server/src/test/java/vip/mate/channel/web/ChatStreamTrackerOrphanPolicyTest.java @@ -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 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 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 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 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 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 runs = + (Map) 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 runs = + (Map) 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)); + } } diff --git a/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatApprovalInteractionTest.java b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatApprovalInteractionTest.java index f2cb57d4..62d7e65f 100644 --- a/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatApprovalInteractionTest.java +++ b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatApprovalInteractionTest.java @@ -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 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"); + } } diff --git a/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatDisconnectTestSupport.java b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatDisconnectTestSupport.java new file mode 100644 index 00000000..1e72f478 --- /dev/null +++ b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatDisconnectTestSupport.java @@ -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 tasks = new ArrayDeque<>(); + private boolean shutdown; + + @Override + public void shutdown() { + shutdown = true; + } + + @Override + public List 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(); + } + } +} diff --git a/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatStreamE2ETest.java b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatStreamE2ETest.java index 9453abbc..6c548dbf 100644 --- a/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatStreamE2ETest.java +++ b/mateclaw-server/src/test/java/vip/mate/channel/webchat/WebChatStreamE2ETest.java @@ -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 {