diff --git a/ui/src/api/client.ts b/ui/src/api/client.ts index c061b0d1b6..7723806c9e 100644 --- a/ui/src/api/client.ts +++ b/ui/src/api/client.ts @@ -191,6 +191,9 @@ export interface AgentInstancesApi { ): Promise; get(id: string, options?: ReadOptions): Promise; + /** Restores a crashed runtime to its last retained snapshot. */ + revert(id: string): Promise; + create(input: { agent: ResourceRefInput; /** @@ -355,6 +358,7 @@ export function createApiClient(): KagentApiClient { invoke("agentInstances.rename", { id, name }), get: (id, options) => invoke("agentInstances.get", { id }, options), + revert: (id) => invoke("agentInstances.revert", { id }), create: (input) => invoke("agentInstances.create", input), remove: (id) => invoke("agentInstances.delete", { id }), fork: (id, name) => diff --git a/ui/src/api/domain/agentInstances.ts b/ui/src/api/domain/agentInstances.ts index 68568298fe..76bdd51b9b 100644 --- a/ui/src/api/domain/agentInstances.ts +++ b/ui/src/api/domain/agentInstances.ts @@ -58,6 +58,17 @@ export type AgentInstanceOperation = | "delete" | "unknown"; +export interface SessionMessageReference { + taskId: string; + messageId: string; +} + +export interface AgentInstanceRecovery { + /** Absent when the controller could not observe the runtime. */ + needed?: boolean; + failedMessages: SessionMessageReference[]; +} + /** * Why an instance failed, as the controller recorded it. * @@ -106,6 +117,7 @@ export interface AgentInstance { operation: AgentInstanceOperation; /** Set only when the controller recorded a failure. */ failure?: AgentInstanceFailure; + recovery?: AgentInstanceRecovery; /** RFC3339, or empty when the record carried no timestamp. */ createdAt: string; /** RFC3339, or empty when the record carried no timestamp. */ diff --git a/ui/src/api/grpc/operations.ts b/ui/src/api/grpc/operations.ts index 3fd6ded89b..ef9e888fed 100644 --- a/ui/src/api/grpc/operations.ts +++ b/ui/src/api/grpc/operations.ts @@ -594,6 +594,15 @@ function toAgentInstance(instance: PbSession): AgentInstance { message: orUndefined(instance.failure.message), } : undefined, + recovery: instance.recovery + ? { + needed: instance.recovery.needed, + failedMessages: instance.recovery.failedMessages.map((message) => ({ + taskId: message.taskId, + messageId: message.messageId, + })), + } + : undefined, createdAt: isoFrom(instance.createdAt), updatedAt: isoFrom(instance.updatedAt), }; @@ -666,6 +675,7 @@ const agentInstances: Pick< | "agentInstances.rename" | "agentInstances.fork" | "agentInstances.delete" + | "agentInstances.revert" > = { "agentInstances.list": async (input, options) => { const name = "SessionService/ListSessions"; @@ -941,6 +951,16 @@ const agentInstances: Pick< return toAgentInstance(required(response.session, name, "agent instance")); }, + "agentInstances.revert": async (input, options) => { + const name = "SessionService/RevertSession"; + const response = await rpc(name, options.signal, () => + serviceClient(SessionService).revertSession( + { sessionId: input.id }, + call("agentInstances.revert", options), + ), + ); + return toAgentInstance(required(response.session, name, "agent instance")); + }, }; // endregion diff --git a/ui/src/api/hooks/useLiveTranscript.ts b/ui/src/api/hooks/useLiveTranscript.ts index 1c04604848..a290e6a18d 100644 --- a/ui/src/api/hooks/useLiveTranscript.ts +++ b/ui/src/api/hooks/useLiveTranscript.ts @@ -48,3 +48,17 @@ export function useLiveTranscript( }; }, [enabled, isBusy]); } + +/** + * Re-reads a conversation's transcript and its Session record together. + * + * Settled rather than awaited as one: a failed Session read must not drop the + * transcript refresh, and the reverse. Each keeps its previous reading on failure + * and the next tick tries again, so neither rejection escapes as unhandled. + */ +export async function refreshConversation( + transcript: () => Promise, + session: () => Promise, +): Promise { + await Promise.allSettled([transcript(), session()]); +} diff --git a/ui/src/api/operations.test.ts b/ui/src/api/operations.test.ts index b862f8b9c7..5f9c866222 100644 --- a/ui/src/api/operations.test.ts +++ b/ui/src/api/operations.test.ts @@ -994,6 +994,40 @@ describe("agent instances", () => { expect(isNotFound(missing)).toBe(true); }); + it("reverts through the existing Session RPC and retains exact failed message references", async () => { + let received: string | undefined; + serve(({ service }) => { + service(SessionService, { + revertSession: (request) => { + received = request.sessionId; + return { session: instanceMessage({ + state: RuntimeState.SUSPENDED, + recovery: { needed: false, failedMessages: [ + { taskId: "completed-task", messageId: "discarded-message" }, + ] }, + }) }; + }, + }); + }); + const reverted = await apiClient.agentInstances.revert(INSTANCE_ID); + expect(received).toBe(INSTANCE_ID); + expect(reverted.state).toBe("suspended"); + expect(reverted.recovery).toEqual({ needed: false, failedMessages: [ + { taskId: "completed-task", messageId: "discarded-message" }, + ] }); + }); + + it("keeps an unknown recovery observation distinct from not needing recovery", async () => { + serve(({ service }) => { + service(SessionService, { + getSession: () => ({ session: instanceMessage({ recovery: {} }) }), + }); + }); + expect((await apiClient.agentInstances.get(INSTANCE_ID)).recovery).toEqual({ + needed: undefined, failedMessages: [], + }); + }); + }); describe("session shares", () => { diff --git a/ui/src/api/operations.ts b/ui/src/api/operations.ts index 99278d5f49..940f3cbb5f 100644 --- a/ui/src/api/operations.ts +++ b/ui/src/api/operations.ts @@ -288,6 +288,7 @@ export interface OperationMap { output: void; }; + "agentInstances.revert": { input: AgentInstanceRef; output: AgentInstance }; "harnesses.list": { input: { namespace?: string }; output: Harness[] }; /** diff --git a/ui/src/components/chat/ChatMessageItem.tsx b/ui/src/components/chat/ChatMessageItem.tsx index 8e9a0f28df..7a57e0cecb 100644 --- a/ui/src/components/chat/ChatMessageItem.tsx +++ b/ui/src/components/chat/ChatMessageItem.tsx @@ -12,6 +12,7 @@ import { ToolApprovalRecord } from "./ToolApprovalRecord"; import { AskUserRecord } from "./AskUserRecord"; import { AttachmentChip } from "./AttachmentChip"; import { isAwaitingContent, messageText } from "./messageText"; +import type { DiscardedMark } from "./discardedTurns"; const { Text } = Typography; @@ -29,6 +30,7 @@ export function ChatMessageItem({ message, sessionId, isCheckpointed = false, + discarded, onManagePush, }: { message: ChatMessage; @@ -36,6 +38,8 @@ export function ChatMessageItem({ sessionId?: string; /** Whether this message is above the nearest saved boundary, for the browser suite. */ isCheckpointed?: boolean; + /** Whether a runtime recovery discarded this message, and how; see `discardedMessages`. */ + discarded?: DiscardedMark; /** Opens the callback settings for this message's task. */ onManagePush?: () => void; }) { @@ -54,6 +58,7 @@ export function ChatMessageItem({ data-message-id={message.id} data-role={message.role} data-checkpointed={isCheckpointed || undefined} + data-discarded={discarded} css={{ display: "grid", gap: theme.space(2), @@ -72,6 +77,13 @@ export function ChatMessageItem({ {isUser ? "You" : "Agent"} + {/* Labelled in the header rather than only dimmed: the reader message above it + carries the explanation, and a dimmed reply on its own reads as styling. */} + {discarded === "reply" ? ( + + ยท Not remembered after recovery + + ) : null} {onManagePush ? ( + + )} + /> + ); +} diff --git a/ui/src/components/chat/discardedTurns.ts b/ui/src/components/chat/discardedTurns.ts new file mode 100644 index 0000000000..0d85369f24 --- /dev/null +++ b/ui/src/components/chat/discardedTurns.ts @@ -0,0 +1,44 @@ +import type { AgentInstanceRecovery, ChatMessage } from "@/api"; + +/** + * How a message relates to work a runtime recovery discarded. + * + * - `message`: a reader message the controller listed as discarded. It did reach the + * agent, which may have started answering, so it is not a message that failed to send. + * - `reply`: anything after that message in the same task. The controller lists reader + * messages only, but the agent output that followed one was produced by the runtime + * state the recovery threw away, so it is unremembered too. + */ +export type DiscardedMark = "message" | "reply"; + +/** + * Marks each message the recovered runtime does not remember, by message id. + * + * A message is matched on both its task and message id, because ids are only unique + * within a task. Marking starts at the first listed message of a task, not at the task: + * a task can hold answers captured in the snapshot before the crash, and those stay + * unmarked. This reads transcript order, so it marks a reply correctly only when the + * transcript places the reply after the message it answers. + */ +export function discardedMessages( + messages: readonly ChatMessage[], + recovery: AgentInstanceRecovery | undefined, +): ReadonlyMap { + const marks = new Map(); + const failed = recovery?.failedMessages ?? []; + if (failed.length === 0) return marks; + + const listed = new Set(failed.map((ref) => `${ref.taskId}\u0000${ref.messageId}`)); + const discardedTasks = new Set(); + for (const message of messages) { + const taskId = message.taskId; + if (!taskId) continue; + if (listed.has(`${taskId}\u0000${message.id}`)) { + marks.set(message.id, "message"); + discardedTasks.add(taskId); + } else if (discardedTasks.has(taskId)) { + marks.set(message.id, "reply"); + } + } + return marks; +} diff --git a/ui/src/mocks/mockBackend.test.ts b/ui/src/mocks/mockBackend.test.ts index 8a079e7ee3..b2589dc0e5 100644 --- a/ui/src/mocks/mockBackend.test.ts +++ b/ui/src/mocks/mockBackend.test.ts @@ -108,6 +108,7 @@ const INPUTS = { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44", }, + "agentInstances.revert": { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44" }, "agentInstances.shares.list": { id: "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44", @@ -246,8 +247,10 @@ describe("the fixture backend", () => { operationIds.map(async (id) => { try { await run(id); + if (id === "agentInstances.revert") return `${id}: a runtime that is not crashed accepted recovery`; return null; } catch (error) { + if (id === "agentInstances.revert" && error instanceof ApiError && error.code === "FailedPrecondition") return null; return `${id}: ${(error as Error).message}`; } }), diff --git a/ui/src/mocks/transport.ts b/ui/src/mocks/transport.ts index 6756b51f1c..f4827ad161 100644 --- a/ui/src/mocks/transport.ts +++ b/ui/src/mocks/transport.ts @@ -644,6 +644,7 @@ function agentInstanceMessage( // turns it back into "not reported". preparedRevision: row.preparedRevision ?? "", a2aAuthority: row.a2aAuthority ?? "", + recovery: row.recovery, state: PB_STATE_BY_NAME[row.state], operation: PB_OPERATION_BY_NAME[row.operation], failure: row.failure @@ -825,6 +826,12 @@ on(SessionService.method.getSession, (input, call) => ({ ), })); +// Fixtures have no Substrate runtime whose crashed state can authorize a revert. +on(SessionService.method.revertSession, (input, call) => { + instanceFor(requireInstanceId(input.sessionId), call); + throw new ConnectError("The fixture runtime is not crashed", Code.FailedPrecondition); +}); + on(SessionService.method.createSession, (input, call) => { const namespace = requireNamespace(input.agent?.namespace ?? ""); if (!input.agent?.name.trim()) throw new ConnectError("an Agent is required", Code.InvalidArgument); diff --git a/ui/src/pages/AgentChatPage.recovery.test.tsx b/ui/src/pages/AgentChatPage.recovery.test.tsx new file mode 100644 index 0000000000..4639c5c646 --- /dev/null +++ b/ui/src/pages/AgentChatPage.recovery.test.tsx @@ -0,0 +1,156 @@ +import { act, render, screen, waitFor } from "@testing-library/react"; +import userEvent from "@testing-library/user-event"; +import { ThemeProvider } from "@emotion/react"; +import { MemoryRouter, Route, Routes } from "react-router-dom"; +import { SWRConfig } from "swr"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { createRouterTransport } from "@connectrpc/connect"; +import { SessionService } from "@/generated/kagent/api/v1/sessions_pb"; +import { CheckpointService } from "@/generated/kagent/api/v1/checkpoints_pb"; +import { RuntimeOperation, RuntimeState } from "@/generated/kagent/api/v1/runtime_pb"; +import { setApiTransport } from "@/api/transport"; +import type { ChatController } from "@/api"; +import { themeFor } from "@/theme/theme"; +import { AgentChatPage } from "./AgentChatPage"; + +/* + * The page, against a fake controller whose Session can change under it. + * + * Only the chat controller is replaced: it speaks A2A, which is not what these tests + * are about. Every Session read goes through the real hooks and the real client, so a + * test that passes here means the page itself discovers the crash and reacts to it. + */ +const chat = { + messages: [], + isLoadingHistory: false, + turnState: "idle", + turnPhase: "idle", + phase: "idle", + refreshTranscript: vi.fn(() => Promise.resolve()), + send: vi.fn(), + cancel: vi.fn(), + dismissQuestion: vi.fn(), + answerQuestion: vi.fn(), + answerToolApproval: vi.fn(), + retry: vi.fn(), +} as unknown as ChatController; + +vi.mock("@/api", async (importOriginal) => ({ + ...(await importOriginal()), + useChat: () => chat, +})); + +const ID = "6f1c9d20-1b7a-4a1e-9a3f-2c0d8e5b1a44"; +const LIVE_REFRESH_MS = 4_000; + +type Recovery = { needed?: boolean; failedMessages: { taskId: string; messageId: string }[] }; + +function session(recovery?: Recovery) { + return { + id: ID, + contextId: ID, + creator: "alice@example.com", + agent: { ref: { namespace: "kagent", name: "assistant" }, uid: "agent-uid" }, + state: RuntimeState.READY, + operation: RuntimeOperation.NONE, + recovery, + }; +} + +const HEALTHY: Recovery = { needed: false, failedMessages: [] }; +const CRASHED: Recovery = { needed: true, failedMessages: [] }; + +let backend: Recovery; +let getSession: ReturnType; +let revertSession: ReturnType; + +function renderPage() { + setApiTransport(createRouterTransport(({ service }) => { + service(SessionService, { + getSession: (getSession = vi.fn(() => ({ session: session(backend) }))), + listSessions: () => ({ sessions: [session(backend)] }), + revertSession: (revertSession = vi.fn(() => { + backend = HEALTHY; + return { session: session(backend) }; + })), + }); + service(CheckpointService, { listCheckpoints: () => ({ checkpoints: [] }) }); + })); + render( + + new Map(), shouldRetryOnError: false, revalidateOnFocus: false }}> + + + } /> + + + + , + ); +} + +const composer = () => screen.getByTestId("chat-input"); + +/** Fires interval ticks, then lets the reads they started settle. */ +async function tick(ms: number) { + await act(async () => { + vi.advanceTimersByTime(ms); + }); +} + +beforeEach(() => { + // Only intervals are faked: the page's polls are intervals, and SWR and the + // transport keep their real timeouts and promises. + vi.useFakeTimers({ toFake: ["setInterval", "clearInterval"] }); + vi.mocked(chat.refreshTranscript).mockClear(); +}); + +afterEach(() => { + vi.useRealTimers(); + setApiTransport(undefined); +}); + +describe("a crash found on an open conversation", () => { + it("finds a crash that happens after the page went idle", async () => { + backend = HEALTHY; + renderPage(); + await waitFor(() => expect(composer()).toBeEnabled()); + expect(screen.queryByText("The agent runtime crashed")).not.toBeInTheDocument(); + + // The Actor crashes with no turn running, so nothing on the page is polling closely. + backend = CRASHED; + await tick(LIVE_REFRESH_MS); + + expect(await screen.findByText("The agent runtime crashed")).toBeInTheDocument(); + expect(composer()).toBeDisabled(); + }); + + it("does not re-read a crashed Session every second while it waits on the reader", async () => { + backend = CRASHED; + renderPage(); + await screen.findByText("The agent runtime crashed"); + await waitFor(() => expect(composer()).toBeDisabled()); + + const before = getSession.mock.calls.length; + await tick(LIVE_REFRESH_MS - 1); + expect(getSession.mock.calls.length).toBe(before); + await tick(1); + await waitFor(() => expect(getSession.mock.calls.length).toBe(before + 1)); + }); + + it("reverts, then re-reads the Session and the transcript and enables the composer", async () => { + backend = CRASHED; + renderPage(); + await screen.findByText("The agent runtime crashed"); + expect(composer()).toBeDisabled(); + const reads = getSession.mock.calls.length; + + await userEvent.click(screen.getByRole("button", { name: "Revert to last snapshot" })); + + await waitFor(() => expect(composer()).toBeEnabled()); + expect(revertSession).toHaveBeenCalledOnce(); + expect(getSession.mock.calls.length).toBeGreaterThan(reads); + expect(chat.refreshTranscript).toHaveBeenCalled(); + expect(screen.queryByText("The agent runtime crashed")).not.toBeInTheDocument(); + }); +}); diff --git a/ui/src/pages/AgentChatPage.tsx b/ui/src/pages/AgentChatPage.tsx index e2241495de..f8f13f69d0 100644 --- a/ui/src/pages/AgentChatPage.tsx +++ b/ui/src/pages/AgentChatPage.tsx @@ -13,6 +13,7 @@ import { ConversationDetailsModal } from "@/components/chat/ConversationDetailsM import { SnapshotDetailsModal } from "@/components/chat/SnapshotDetailsModal"; import { SnapshotRenameDialog } from "@/components/chat/SnapshotRenameDialog"; import { ChatTranscript } from "@/components/chat/ChatTranscript"; +import { SessionRecoveryAlert } from "@/components/chat/SessionRecoveryAlert"; import { TaskPushNotificationsModal } from "@/components/chat/TaskPushNotificationsModal"; import { isLifecycleBusy } from "@/components/chat/lifecycleReading"; import { paths } from "@/router/routes"; @@ -24,7 +25,7 @@ import { useChat, } from "@/api"; import { autoTitleFrom } from "@/components/agent-instances/instanceLabels"; -import { useLiveTranscript } from "@/api/hooks/useLiveTranscript"; +import { refreshConversation, useLiveTranscript } from "@/api/hooks/useLiveTranscript"; import { useInvalidateConversations } from "@/api/hooks/useInvalidateConversations"; import { useCheckpoints } from "@/api/hooks/useCheckpoints"; import type { Checkpoint } from "@/api"; @@ -114,8 +115,17 @@ export function AgentChatPage() { const chat = useChat(conversation); const [pushTask, setPushTask] = useState<{ conversationId: string; taskId: string }>(); - // Keep up with messages sent from other tabs or clients. - useLiveTranscript(chat.refreshTranscript, { + /* + * Keep up with messages sent from other tabs or clients, and with the Session itself. + * + * The Session is re-read on the same tick because a runtime can crash while nothing + * here is watching: a turn settles, the lifecycle poll below stops, and the Actor + * crashes during automatic suspension. A transcript-only refresh would never learn + * that, so the composer would stay enabled and the recovery choices hidden until the + * reader reloaded. Riding this refresh costs one read per tick while the tab is + * visible, rather than widening the one-second lifecycle poll to every idle page. + */ + useLiveTranscript(() => refreshConversation(chat.refreshTranscript, instance.refresh), { enabled: Boolean(conversation), isBusy: chat.phase === "streaming", }); @@ -305,7 +315,8 @@ export function AgentChatPage() { // Automatic runtime quiescence leaves the conversation ready for its next turn. const state = instance.data?.state; - const canSend = state === "ready"; + const recovery = instance.data?.recovery; + const canSend = state === "ready" && recovery?.needed !== true; /* * Watching the instance while it is doing something. @@ -427,7 +438,36 @@ export function AgentChatPage() { // and depending on it would re-run this on each one. // eslint-disable-next-line react-hooks/exhaustive-deps }, [id, chat.isLoadingHistory]); - const isBusy = isLifecycleBusy(instance.data, chat.turnPhase, isAwaitingReply); + const [isReverting, setReverting] = useState(false); + const revertRuntime = useCallback(async () => { + if (!id) return; + setReverting(true); + let reverted = false; + try { + await apiClient.agentInstances.revert(id); + reverted = true; + await instance.refresh(); + await chat.refreshTranscript(); + } catch (cause: unknown) { + const reason = cause instanceof Error ? cause.message : String(cause); + toast.error(reverted + ? `Recovery completed. Could not refresh the conversation: ${reason}` + : `Could not recover the conversation: ${reason}`); + } finally { + setReverting(false); + } + }, [id, instance, chat]); + + /* + * A crashed conversation waiting on the reader is idle, so `recovery.needed` alone does + * not keep the one-second poll running: every read of a crashed Session also asks + * Substrate for its Actor, and a page left open would do that until someone acted. + * The poll runs while a revert is in flight or the controller has claimed a resume, + * which `isLifecycleBusy` already reads. The live refresh above still notices a + * recovery made from another tab or by the controller. + */ + const isBusy = isLifecycleBusy(instance.data, chat.turnPhase, isAwaitingReply) + || isReverting; /* * The conversation list is deliberately not re-read when this one's state changes. @@ -620,7 +660,14 @@ export function AgentChatPage() { /> ) : null} - {instance.data && !canSend ? ( + + + {instance.data && !canSend && recovery?.needed !== true ? ( refreshConversation(chat.refreshTranscript, instance.refresh), { enabled: Boolean(conversation), isBusy: chat.phase === "streaming", }); @@ -92,7 +93,7 @@ export function SharedAgentPage() { height: `calc(100vh - ${theme.layout.headerHeight}px - ${theme.space(48)})`, }} > - + )}