Skip to content

Commit 5ae5c00

Browse files
adalundheclaude
adalundhe
and
claude
committed
fix(ui): keep revert usable after an automatic recovery attempt fails
The recovery alert disabled Revert whenever the Session's operation was RESUME. An admitted crash recovery keeps that operation through every failed attempt, so once the first automatic attempt failed the button spun until a later attempt succeeded. That is the moment the reader needs it, and Revert was the only thing to press there. Revert is now disabled only while this page's own RevertSession call is in flight. An admitted recovery still reads "Recovering this conversation", in automatic mode and in manual mode when another tab or client admitted it. A click joins the admitted recovery, as the service does for an admitted revert. A ref guards a double click before the button re-renders disabled. An ABORTED refusal from other lifecycle work shows its error and leaves the button usable. A FAILED_PRECONDITION with reason KAGENT_SESSION_RUNTIME_NOT_CRASHED means another attempt recovered it first, so the page re-reads the Session and transcript without an error. Leaving Revert enabled while an attempt runs was preferred to polling for the attempt's outcome: the service already joins or refuses a concurrent revert, so the button needs no second copy of that rule. Verified with Node 24 on Linux: tsc --noEmit is clean, eslint has no errors, and vitest run src/components/chat src/pages src/api passes 379 tests. The six new tests fail against the previous alert and page. On a kind fleet with automatic recovery on, the automatic attempt was held failing; Revert stayed enabled and restored the Session. Signed-off-by: Ada Lundhe <[email protected]> Co-Authored-By: Claude Opus 5.5 <[email protected]>
1 parent 3fc96de commit 5ae5c00

4 files changed

Lines changed: 166 additions & 15 deletions

File tree

‎ui/src/components/chat/SessionRecoveryAlert.test.tsx‎

Lines changed: 34 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -9,11 +9,11 @@ import { SessionRecoveryAlert } from "./SessionRecoveryAlert";
99
const ready: Checkpoint = { id: "ready", name: "Saved conversation", agentInstanceId: "session", headTaskId: "task", state: "ready" };
1010
const failed: Checkpoint = { id: "failed", name: "Not saved", agentInstanceId: "session", headTaskId: "task", state: "failed" };
1111

12-
function show(recovery?: AgentInstanceRecovery, { isRecovering = false, checkpoints = [ready] } = {}) {
12+
function show(recovery?: AgentInstanceRecovery, { isRecovering = false, isReverting = false, checkpoints = [ready] } = {}) {
1313
const onRevert = vi.fn();
1414
render(
1515
<ThemeProvider theme={themeFor("dark")}>
16-
<SessionRecoveryAlert {...{ recovery, checkpoints, isRecovering, onRevert }} />
16+
<SessionRecoveryAlert {...{ recovery, checkpoints, isRecovering, isReverting, onRevert }} />
1717
</ThemeProvider>,
1818
);
1919
return { onRevert };
@@ -63,9 +63,39 @@ describe("conversation recovery choices", () => {
6363
expect(screen.queryByText(/marked as failed/)).not.toBeInTheDocument();
6464
});
6565

66-
it("prevents another manual request while recovery is running", () => {
67-
show({ needed: true, automatic: false, failedMessages: [] }, { isRecovering: true });
66+
/*
67+
* An admitted recovery keeps its operation through every failed attempt. Disabling
68+
* Revert for as long as one was admitted left the button spinning once the first
69+
* automatic attempt failed, with nothing else to press.
70+
*/
71+
it.each([
72+
{ mode: "automatic", automatic: true },
73+
{ mode: "manual, admitted from another tab", automatic: false },
74+
])("keeps Revert usable while an admitted recovery is pending ($mode)", async ({ automatic }) => {
75+
const { onRevert } = show({ needed: true, automatic, failedMessages: [] }, { isRecovering: true });
76+
expect(screen.getByText("Recovering this conversation")).toBeInTheDocument();
77+
const revert = screen.getByRole("button", { name: "Revert to last snapshot" });
78+
expect(revert).toBeEnabled();
79+
await userEvent.click(revert);
80+
expect(onRevert).toHaveBeenCalledOnce();
81+
});
82+
83+
it("disables Revert only while this page's revert is in flight", () => {
84+
const recovery: AgentInstanceRecovery = { needed: true, automatic: true, failedMessages: [] };
85+
const props = { recovery, checkpoints: [ready], isRecovering: true, onRevert: vi.fn() };
86+
const { rerender } = render(
87+
<ThemeProvider theme={themeFor("dark")}>
88+
<SessionRecoveryAlert {...props} isReverting />
89+
</ThemeProvider>,
90+
);
6891
expect(screen.getByRole("button", { name: /Revert to last snapshot/ })).toBeDisabled();
92+
// The request failed; the admitted recovery is still pending, and so is the choice.
93+
rerender(
94+
<ThemeProvider theme={themeFor("dark")}>
95+
<SessionRecoveryAlert {...props} isReverting={false} />
96+
</ThemeProvider>,
97+
);
98+
expect(screen.getByRole("button", { name: /Revert to last snapshot/ })).toBeEnabled();
6999
expect(screen.getByText("Recovering this conversation")).toBeInTheDocument();
70100
});
71101
});

‎ui/src/components/chat/SessionRecoveryAlert.tsx‎

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,28 @@ import { canForkFrom, type AgentInstanceRecovery, type Checkpoint } from "@/api"
1414
* choices rather than only showing progress. A revert joins the recovery already
1515
* admitted instead of starting a second one, and a fork reads only a saved snapshot, so
1616
* neither races the automatic attempt.
17+
*
18+
* That is why `isRecovering` and `isReverting` are separate. An admitted recovery keeps
19+
* its operation until a revert succeeds, so it outlives every failed attempt. Disabling
20+
* Revert for as long as one was admitted left the button spinning after the first
21+
* automatic attempt failed, which is the one moment the reader needs it. Only this
22+
* page's own request in flight disables it; a click while another attempt runs joins
23+
* that attempt, and one refused because other lifecycle work holds the runtime comes
24+
* back as an error the reader can retry.
1725
*/
1826
export function SessionRecoveryAlert({
1927
recovery,
2028
checkpoints,
2129
isRecovering,
30+
isReverting,
2231
onRevert,
2332
}: {
2433
recovery?: AgentInstanceRecovery;
2534
checkpoints: readonly Checkpoint[];
35+
/** A recovery is admitted for this conversation, by anyone, or this page is reverting. */
2636
isRecovering: boolean;
37+
/** This page's own revert request is in flight. */
38+
isReverting: boolean;
2739
onRevert: () => void | Promise<void>;
2840
}) {
2941
if (recovery?.needed !== true) return null;
@@ -46,8 +58,8 @@ export function SessionRecoveryAlert({
4658
{" "}The reverted agent does not remember messages sent after that snapshot.
4759
</div>
4860
<Button
49-
loading={isRecovering}
50-
disabled={isRecovering}
61+
loading={isReverting}
62+
disabled={isReverting}
5163
onClick={() => void onRevert()}
5264
style={{ marginTop: 8 }}
5365
>

‎ui/src/pages/AgentChatPage.recovery.test.tsx‎

Lines changed: 95 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@ import { ThemeProvider } from "@emotion/react";
44
import { MemoryRouter, Route, Routes } from "react-router-dom";
55
import { SWRConfig } from "swr";
66
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
7-
import { createRouterTransport } from "@connectrpc/connect";
7+
import { Code, ConnectError, createRouterTransport } from "@connectrpc/connect";
8+
import toast from "react-hot-toast";
9+
import { ErrorInfoSchema } from "@/generated/google/rpc/error_details_pb";
810
import { SessionService } from "@/generated/kagent/api/v1/sessions_pb";
911
import { CheckpointService } from "@/generated/kagent/api/v1/checkpoints_pb";
1012
import { RuntimeOperation, RuntimeState } from "@/generated/kagent/api/v1/runtime_pb";
@@ -43,36 +45,47 @@ vi.mock("@/api", async (importOriginal) => ({
4345
const ID = "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44";
4446
const LIVE_REFRESH_MS = 4_000;
4547

46-
type Recovery = { needed?: boolean; failedMessages: { taskId: string; messageId: string }[] };
48+
type Recovery = { needed?: boolean; automatic?: boolean; failedMessages: { taskId: string; messageId: string }[] };
49+
50+
/** The Session's lifecycle, apart from its recovery reading. */
51+
let lifecycle: { state: RuntimeState; operation: RuntimeOperation };
4752

4853
function session(recovery?: Recovery) {
4954
return {
5055
id: ID,
5156
contextId: ID,
5257
creator: "[email protected]",
5358
agent: { namespace: "kagent", name: "assistant" },
54-
state: RuntimeState.READY,
55-
operation: RuntimeOperation.NONE,
59+
...lifecycle,
5660
recovery,
5761
};
5862
}
5963

6064
const HEALTHY: Recovery = { needed: false, failedMessages: [] };
6165
const CRASHED: Recovery = { needed: true, failedMessages: [] };
66+
const IDLE = { state: RuntimeState.READY, operation: RuntimeOperation.NONE };
67+
/** An automatic attempt admitted the recovery, then failed; the next one is minutes away. */
68+
const ADMITTED = { state: RuntimeState.FAILED, operation: RuntimeOperation.RESUME };
6269

6370
let backend: Recovery;
6471
let getSession: ReturnType<typeof vi.fn>;
6572
let revertSession: ReturnType<typeof vi.fn>;
73+
function recovered() {
74+
backend = HEALTHY;
75+
lifecycle = IDLE;
76+
return { session: session(backend) };
77+
}
78+
79+
type Reverted = ReturnType<typeof recovered>;
80+
/** What RevertSession does; by default it recovers the Session. */
81+
let revert: () => Reverted | Promise<Reverted>;
6682

6783
function renderPage() {
6884
setApiTransport(createRouterTransport(({ service }) => {
6985
service(SessionService, {
7086
getSession: (getSession = vi.fn(() => ({ session: session(backend) }))),
7187
listSessions: () => ({ sessions: [session(backend)] }),
72-
revertSession: (revertSession = vi.fn(() => {
73-
backend = HEALTHY;
74-
return { session: session(backend) };
75-
})),
88+
revertSession: (revertSession = vi.fn(() => revert())),
7689
});
7790
service(CheckpointService, { listCheckpoints: () => ({ checkpoints: [] }) });
7891
}));
@@ -103,9 +116,12 @@ beforeEach(() => {
103116
// transport keep their real timeouts and promises.
104117
vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] });
105118
vi.mocked(chat.refreshTranscript).mockClear();
119+
lifecycle = IDLE;
120+
revert = recovered;
106121
});
107122

108123
afterEach(() => {
124+
vi.restoreAllMocks();
109125
vi.useRealTimers();
110126
setApiTransport(undefined);
111127
});
@@ -154,3 +170,74 @@ describe("a crash found on an open conversation", () => {
154170
expect(screen.queryByText("The agent runtime crashed")).not.toBeInTheDocument();
155171
});
156172
});
173+
174+
/*
175+
* An admitted recovery keeps its RESUME operation through every failed attempt. The page
176+
* used to disable Revert for as long as that operation stood, so once the first automatic
177+
* attempt failed the button spun until a later attempt succeeded, with nothing to press.
178+
*/
179+
describe("a recovery admitted by a failed automatic attempt", () => {
180+
const revertButton = () => screen.getByRole("button", { name: /Revert to last snapshot/ });
181+
182+
it("keeps Revert usable, and a click joins the admitted recovery", async () => {
183+
backend = { needed: true, automatic: true, failedMessages: [] };
184+
lifecycle = ADMITTED;
185+
renderPage();
186+
await screen.findByText("Recovering this conversation");
187+
await waitFor(() => expect(revertButton()).toBeEnabled());
188+
189+
await userEvent.click(revertButton());
190+
191+
await waitFor(() => expect(screen.queryByText("Recovering this conversation")).not.toBeInTheDocument());
192+
expect(revertSession).toHaveBeenCalledOnce();
193+
expect(chat.refreshTranscript).toHaveBeenCalled();
194+
});
195+
196+
it("sends one request for a double click, and leaves Revert usable when other work refuses it", async () => {
197+
const error = vi.spyOn(toast, "error");
198+
let refuse!: () => void;
199+
revert = () => new Promise<Reverted>((_, reject) => {
200+
refuse = () => reject(new ConnectError("lifecycle operation is pending", Code.Aborted));
201+
});
202+
backend = { needed: true, automatic: true, failedMessages: [] };
203+
lifecycle = ADMITTED;
204+
renderPage();
205+
await screen.findByText("Recovering this conversation");
206+
await waitFor(() => expect(revertButton()).toBeEnabled());
207+
208+
// Two clicks before React re-renders the button disabled.
209+
act(() => {
210+
revertButton().click();
211+
revertButton().click();
212+
});
213+
await waitFor(() => expect(revertButton()).toBeDisabled());
214+
expect(revertSession).toHaveBeenCalledOnce();
215+
216+
await act(async () => refuse());
217+
218+
await waitFor(() => expect(revertButton()).toBeEnabled());
219+
expect(error).toHaveBeenCalledWith(expect.stringContaining("Could not recover the conversation"));
220+
expect(screen.getByText("Recovering this conversation")).toBeInTheDocument();
221+
});
222+
223+
it("re-reads quietly when another attempt recovered the conversation first", async () => {
224+
const error = vi.spyOn(toast, "error");
225+
revert = () => {
226+
recovered();
227+
throw new ConnectError("runtime is not crashed", Code.FailedPrecondition, undefined, [
228+
{ desc: ErrorInfoSchema, value: { domain: "kagent.dev", reason: "KAGENT_SESSION_RUNTIME_NOT_CRASHED" } },
229+
]);
230+
};
231+
backend = { needed: true, automatic: true, failedMessages: [] };
232+
lifecycle = ADMITTED;
233+
renderPage();
234+
await screen.findByText("Recovering this conversation");
235+
await waitFor(() => expect(revertButton()).toBeEnabled());
236+
237+
await userEvent.click(revertButton());
238+
239+
await waitFor(() => expect(screen.queryByText("Recovering this conversation")).not.toBeInTheDocument());
240+
expect(revertSession).toHaveBeenCalledOnce();
241+
expect(error).not.toHaveBeenCalled();
242+
});
243+
});

‎ui/src/pages/AgentChatPage.tsx‎

Lines changed: 23 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import { ChatTranscript } from "@/components/chat/ChatTranscript";
1616
import { SessionRecoveryAlert } from "@/components/chat/SessionRecoveryAlert";
1717
import { TaskPushNotificationsModal } from "@/components/chat/TaskPushNotificationsModal";
1818
import { isLifecycleBusy } from "@/components/chat/lifecycleReading";
19+
import { isApiError } from "@/api/ApiError";
1920
import { paths } from "@/router/routes";
2021
import {
2122
apiClient,
@@ -438,9 +439,17 @@ export function AgentChatPage() {
438439
// and depending on it would re-run this on each one.
439440
// eslint-disable-next-line react-hooks/exhaustive-deps
440441
}, [id, chat.isLoadingHistory]);
442+
/*
443+
* Revert is disabled only while this page's own request is in flight. An admitted
444+
* recovery is not enough: it keeps its operation through every failed attempt, and a
445+
* revert sent meanwhile joins it rather than starting another. The ref closes the gap
446+
* before React re-renders the disabled button, so a double click sends one request.
447+
*/
441448
const [isReverting, setReverting] = useState(false);
449+
const revertInFlight = useRef(false);
442450
const revertRuntime = useCallback(async () => {
443-
if (!id) return;
451+
if (!id || revertInFlight.current) return;
452+
revertInFlight.current = true;
444453
setReverting(true);
445454
let reverted = false;
446455
try {
@@ -449,11 +458,23 @@ export function AgentChatPage() {
449458
await instance.refresh();
450459
await chat.refreshTranscript();
451460
} catch (cause: unknown) {
461+
if (!reverted && isApiError(cause) && cause.reason === "KAGENT_SESSION_RUNTIME_NOT_CRASHED") {
462+
// Another attempt, automatic or from another tab, recovered it first. The
463+
// conversation is fine; re-reading it clears the alert.
464+
try {
465+
await instance.refresh();
466+
await chat.refreshTranscript();
467+
} catch (refreshCause: unknown) {
468+
console.error("Could not re-read the recovered conversation:", refreshCause);
469+
}
470+
return;
471+
}
452472
const reason = cause instanceof Error ? cause.message : String(cause);
453473
toast.error(reverted
454474
? `Recovery completed. Could not refresh the conversation: ${reason}`
455475
: `Could not recover the conversation: ${reason}`);
456476
} finally {
477+
revertInFlight.current = false;
457478
setReverting(false);
458479
}
459480
}, [id, instance, chat]);
@@ -664,6 +685,7 @@ export function AgentChatPage() {
664685
recovery={recovery}
665686
checkpoints={checkpoints.data ?? []}
666687
isRecovering={isReverting || instance.data?.operation === "resume"}
688+
isReverting={isReverting}
667689
onRevert={revertRuntime}
668690
/>
669691

0 commit comments

Comments
 (0)