wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Preserve running session replay state

Browse Source
This commit is contained in:
武阳
2026-05-07 16:49:42 +08:00
parent 2ba67ead00
commit 61a0a0a648
5 changed files with 213 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
frontend/app/components/assistant-ui/thread-list.tsx
+16 -2
View File
@@ -50,6 +50,7 @@ export type ClawRunStatus = {
session_id?: string;
status?: "idle" | "queued" | "running" | "completed" | "failed" | "cancelled";
current_stage?: string;
started_at?: number;
};
type ClawStoredToolCall = {
@@ -431,7 +432,11 @@ export function toReplayRepository(
id,
role: "assistant",
parts,
metadata: { sessionId: session.session_id ?? fallbackSessionId },
metadata: {
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
runStartedAtMs: normalizeRunTimestampMs(runStatus.started_at),
},
} as UIMessage;
const threadMessage = {
id,
@@ -444,7 +449,11 @@ export function toReplayRepository(
unstable_annotations: [],
unstable_data: [],
steps: [],
custom: { sessionId: session.session_id ?? fallbackSessionId },
custom: {
sessionId: session.session_id ?? fallbackSessionId,
runId: runStatus.run_id,
runStartedAtMs: normalizeRunTimestampMs(runStatus.started_at),
},
},
} as ThreadMessage;
bindExternalStoreMessage(threadMessage, uiMessage);
@@ -463,6 +472,11 @@ function isActiveRunStatus(
return runStatus?.status === "queued" || runStatus?.status === "running";
}
function normalizeRunTimestampMs(value: unknown) {
if (typeof value !== "number" || !Number.isFinite(value)) return undefined;
return value < 10_000_000_000 ? Math.round(value * 1000) : Math.round(value);
}
function findLastAssistantContentIndex(messages: readonly ClawStoredMessage[]) {
for (let index = messages.length - 1; index >= 0; index -= 1) {
const message = messages[index];
frontend/app/components/assistant-ui/thread.tsx
+113 -27
View File
@@ -157,21 +157,11 @@ export const Thread: FC = () => {
function useRefreshReplayedRun() {
const { replaySession } = useClawSessionReplay();
const isRunning = useAuiState((s) => s.thread.isRunning);
const hasReplayedRunMessage = useAuiState((s) =>
s.thread.messages.some((message) => message.id.includes("-run-")),
);
const latestSessionId = useAuiState((s) =>
findLatestMessageMetadataValue(s.thread.messages, "sessionId"),
);
const replayedRun = useReplayedRunState();
useEffect(() => {
if (!isRunning || !hasReplayedRunMessage) return;
const sessionId = normalizeSessionId(
latestSessionId ?? window.localStorage.getItem("claw.activeSessionId"),
);
if (!sessionId) return;
const activeSessionId = sessionId;
if (!replayedRun.active || !replayedRun.sessionId) return;
const activeSessionId = replayedRun.sessionId;
let cancelled = false;
async function refreshIfFinished() {
@@ -198,7 +188,41 @@ function useRefreshReplayedRun() {
cancelled = true;
window.clearInterval(interval);
};
}, [hasReplayedRunMessage, isRunning, latestSessionId, replaySession]);
}, [replayedRun.active, replayedRun.sessionId, replaySession]);
}
function useReplayedRunState() {
return useAuiState((s) => {
const replayRunMessage = s.thread.messages.find((message) =>
isReplayRunMessageId(message.id),
);
const sessionId = normalizeSessionId(
readMetadataValue(replayRunMessage?.metadata, "sessionId") ??
findLatestMessageMetadataValue(s.thread.messages, "sessionId") ??
(typeof window !== "undefined"
? window.localStorage.getItem("claw.activeSessionId")
: null),
);
const runId = normalizeSessionId(
readMetadataValue(replayRunMessage?.metadata, "runId"),
);
return {
active: Boolean(replayRunMessage),
sessionId,
runId,
};
});
}
async function cancelLatestRun(sessionId: string, runId?: string | null) {
await fetch("/api/claw/runs/cancel", {
method: "POST",
headers: { "content-type": "application/json" },
body: JSON.stringify({
session_id: sessionId,
...(runId ? { run_id: runId } : {}),
}),
});
}
const ChatTopActions: FC = () => {
@@ -315,6 +339,7 @@ const ThreadSuggestionItem: FC = () => {
};
const Composer: FC = () => {
const replayedRun = useReplayedRunState();
return (
<ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
<ComposerPrimitive.AttachmentDropzone asChild>
@@ -329,6 +354,7 @@ const Composer: FC = () => {
rows={1}
autoFocus
aria-label="Message input"
disabled={replayedRun.active}
/>
<ComposerAction />
</div>
@@ -338,6 +364,8 @@ const Composer: FC = () => {
};
const ComposerAction: FC = () => {
const runtimeRunning = useAuiState((s) => s.thread.isRunning);
const replayedRun = useReplayedRunState();
return (
<div className="aui-composer-action-wrapper relative flex items-center justify-between">
<div className="flex items-center gap-1">
@@ -345,7 +373,7 @@ const ComposerAction: FC = () => {
<ComposerAssistButtons />
</div>
<ComposerContextStatus />
<AuiIf condition={(s) => !s.thread.isRunning}>
{!runtimeRunning && !replayedRun.active ? (
<ComposerPrimitive.Send asChild>
<TooltipIconButton
tooltip="Send message"
@@ -359,8 +387,8 @@ const ComposerAction: FC = () => {
<ArrowUpIcon className="aui-composer-send-icon size-4" />
</TooltipIconButton>
</ComposerPrimitive.Send>
</AuiIf>
<AuiIf condition={(s) => s.thread.isRunning}>
) : null}
{runtimeRunning ? (
<ComposerPrimitive.Cancel asChild>
<Button
type="button"
@@ -372,11 +400,58 @@ const ComposerAction: FC = () => {
<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
</Button>
</ComposerPrimitive.Cancel>
</AuiIf>
) : null}
{!runtimeRunning && replayedRun.active ? (
<ReplayedRunCancelButton
sessionId={replayedRun.sessionId}
runId={replayedRun.runId}
/>
) : null}
</div>
);
};
const ReplayedRunCancelButton: FC<{
sessionId: string | null;
runId: string | null;
}> = ({ sessionId, runId }) => {
const { replaySession } = useClawSessionReplay();
const [cancelling, setCancelling] = useState(false);
return (
<Button
type="button"
variant="default"
size="icon"
className="aui-composer-cancel size-8 rounded-full"
aria-label="Stop generating"
disabled={!sessionId || cancelling}
onClick={async () => {
if (!sessionId) return;
setCancelling(true);
try {
await cancelLatestRun(sessionId, runId);
const runStatus = await fetchLatestRunStatus(sessionId);
const response = await fetch(`/api/claw/sessions/${sessionId}`, {
cache: "no-store",
});
if (response.ok) {
const payload = await response.json();
replaySession(
sessionId,
toReplayRepository(payload, sessionId, runStatus),
);
}
window.dispatchEvent(new Event("claw-sessions-changed"));
} finally {
setCancelling(false);
}
}}
>
<SquareIcon className="aui-composer-cancel-icon size-3 fill-current" />
</Button>
);
};
const ComposerContextStatus: FC = () => {
const status = useComposerContextStatus();
const totalTokens = status.contextBudget?.projected_input_tokens ?? 0;
@@ -1130,27 +1205,34 @@ const ActivityChainSummary: FC<{
const { openItem } = useActivityPanel();
const startedAtRef = useRef(Date.now());
const [elapsedMs, setElapsedMs] = useState(0);
const runStartedAtMs = useAuiState((s) => {
const value = readMetadataValue(s.message.metadata, "runStartedAtMs");
return typeof value === "number" && Number.isFinite(value) ? value : null;
});
const label = useAuiState((s) => {
const message = s.message;
return summarizeActivityChainLabel(
message.content,
message.status?.type,
elapsedMs,
);
const status =
message.status?.type === "running" || isReplayRunMessageId(messageId)
? "running"
: message.status?.type;
return summarizeActivityChainLabel(message.content, status, elapsedMs);
});
const active = useAuiState((s) => {
const message = s.message;
return message.status?.type === "running";
return (
message.status?.type === "running" || isReplayRunMessageId(messageId)
);
});
useEffect(() => {
if (!active) return;
setElapsedMs(Date.now() - startedAtRef.current);
const startedAt = runStartedAtMs ?? startedAtRef.current;
setElapsedMs(Math.max(0, Date.now() - startedAt));
const timer = window.setInterval(() => {
setElapsedMs(Date.now() - startedAtRef.current);
setElapsedMs(Math.max(0, Date.now() - startedAt));
}, 1000);
return () => window.clearInterval(timer);
}, [active]);
}, [active, runStartedAtMs]);
return (
<button
@@ -1167,6 +1249,10 @@ const ActivityChainSummary: FC<{
);
};
function isReplayRunMessageId(messageId: string) {
return messageId.includes("-run-");
}
type ActivitySummaryPart = {
type: string;
text?: string;