From c3caa4a1e3cd39bd4f401eba58d17e951e000e2d Mon Sep 17 00:00:00 2001
From: wuyang6 <wuyang6@xiaomi.com>
Date: Fri, 12 Jun 2026 20:07:42 +0800
Subject: [PATCH] Keep optimistic user message during replay

---
 frontend/app/app/assistant.tsx         | 74 +++++++++++++++++++++++---
 frontend/app/lib/claw-session-cache.ts |  8 +++
 2 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/frontend/app/app/assistant.tsx b/frontend/app/app/assistant.tsx
index 3d20005..75b7bf9 100644
--- a/frontend/app/app/assistant.tsx
+++ b/frontend/app/app/assistant.tsx
@@ -8,7 +8,14 @@ import {
 } from "@assistant-ui/react-ai-sdk";
 import type { UIMessage } from "ai";
 import { PanelRightOpenIcon, XIcon } from "lucide-react";
-import { useCallback, useEffect, useMemo, useRef, useState } from "react";
+import {
+	type MutableRefObject,
+	useCallback,
+	useEffect,
+	useMemo,
+	useRef,
+	useState,
+} from "react";
 import {
 	ActivityPanel,
 	ActivityProvider,
@@ -38,6 +45,7 @@ import {
 import {
 	fetchSessionReplaySnapshot,
 	readCachedSessionReplay,
+	type SessionReplaySnapshot,
 } from "@/lib/claw-session-cache";
 import { ClawSessionReplayProvider } from "@/lib/claw-session-replay";
 import { pushSessionUrl, readSessionIdFromUrl } from "@/lib/claw-session-url";
@@ -54,9 +62,12 @@ type AssistantProps = {
 	initialSessionId?: string;
 };
 
+const OPTIMISTIC_REPLAY_GUARD_MS = 30_000;
+
 export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 	const { account, isLoading, setAccount } = useClawAccount();
 	const loadedInitialSessionRef = useRef<string | null>(null);
+	const optimisticSendRef = useRef<Map<string, number>>(new Map());
 	const transport = useMemo(
 		() =>
 			new AssistantChatTransport({
@@ -78,6 +89,9 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 					const lastUserMessage = [...messages]
 						.reverse()
 						.find((message) => message.role === "user");
+					if (lastUserMessage) {
+						optimisticSendRef.current.set(outgoingSessionId, Date.now());
+					}
 					writeActiveSessionId(outgoingSessionId);
 					pushSessionUrl(outgoingSessionId);
 
@@ -126,14 +140,22 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 					Awaited<ReturnType<typeof fetchLatestRunStatus>>
 				>(targetSessionId);
 				if (cached && !cancelled) {
-					replaySession(
-						targetSessionId,
-						toReplayRepository(
-							cached.session,
+					if (
+						!shouldSkipOptimisticEmptyReplay(
+							optimisticSendRef,
 							targetSessionId,
-							cached.runStatus,
-						),
-					);
+							cached,
+						)
+					) {
+						replaySession(
+							targetSessionId,
+							toReplayRepository(
+								cached.session,
+								targetSessionId,
+								cached.runStatus,
+							),
+						);
+					}
 				}
 				const snapshot = await fetchSessionReplaySnapshot<
 					Parameters<typeof toReplayRepository>[0],
@@ -141,6 +163,15 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 				>(targetSessionId, { signal: abortController.signal });
 				if (cancelled) return;
 				if (!snapshot) return;
+				if (
+					shouldSkipOptimisticEmptyReplay(
+						optimisticSendRef,
+						targetSessionId,
+						snapshot,
+					)
+				) {
+					return;
+				}
 				if (!cached || snapshot.signature !== cached.signature) {
 					replaySession(
 						targetSessionId,
@@ -203,6 +234,33 @@ export const Assistant = ({ initialSessionId }: AssistantProps = {}) => {
 	);
 };
 
+function shouldSkipOptimisticEmptyReplay(
+	optimisticSendRef: MutableRefObject<Map<string, number>>,
+	sessionId: string,
+	snapshot: SessionReplaySnapshot | null,
+) {
+	if (hasReplayMessages(snapshot)) {
+		optimisticSendRef.current.delete(sessionId);
+		return false;
+	}
+	const sentAt = optimisticSendRef.current.get(sessionId);
+	if (sentAt === undefined) return false;
+	if (Date.now() - sentAt > OPTIMISTIC_REPLAY_GUARD_MS) {
+		optimisticSendRef.current.delete(sessionId);
+		return false;
+	}
+	return true;
+}
+
+function hasReplayMessages(snapshot: SessionReplaySnapshot | null) {
+	const session = snapshot?.session;
+	if (!session || typeof session !== "object" || Array.isArray(session)) {
+		return false;
+	}
+	const messages = (session as { messages?: unknown }).messages;
+	return Array.isArray(messages) && messages.length > 0;
+}
+
 // 共享 composer 草稿缓存里 newTask（无 sessionId）使用的 key。
 const NEWTASK_KEY = "__newtask__";
 
diff --git a/frontend/app/lib/claw-session-cache.ts b/frontend/app/lib/claw-session-cache.ts
index f63754f..7eac6ff 100644
--- a/frontend/app/lib/claw-session-cache.ts
+++ b/frontend/app/lib/claw-session-cache.ts
@@ -74,6 +74,9 @@ async function fetchSessionStateSnapshot<TSession, TRun>(
 			.sort((left, right) => (left.seq ?? 0) - (right.seq ?? 0))
 			.map((entry) => entry.message)
 			.filter((message) => Boolean(message));
+		if (orderedMessages.length === 0 && !isActiveRunLike(payload.run)) {
+			return null;
+		}
 		const session = {
 			...(payload.session ?? {}),
 			session_id: sessionId,
@@ -157,6 +160,11 @@ function sessionReplaySignature(session: unknown, runStatusValue: unknown) {
 	].join("|");
 }
 
+function isActiveRunLike(value: unknown) {
+	if (!isObject(value)) return false;
+	return value.status === "queued" || value.status === "running";
+}
+
 function isObject(value: unknown): value is Record<string, unknown> {
 	return Boolean(value) && typeof value === "object" && !Array.isArray(value);
 }