From 2d16eaef6b5a8e314f172f782999897713c605df Mon Sep 17 00:00:00 2001 From: sshiv012 <122703005+sshiv012@users.noreply.github.com> Date: Sun, 5 Jul 2026 22:53:14 -0700 Subject: [PATCH] feat(agent): add workflow revert & redo for AI agent turns Per-turn "Revert" rewinds the workflow and the agent-service HEAD to before a turn; a toolbar "Redo" undoes the revert (redo stack, cleared on new prompt). Reuses the ReAct step version-tree + per-step snapshots; adds WsClientRevert/ RedoCommand and WsServerHeadChangeEvent, with revert/redo blocked while the agent is busy and the connect snapshot carrying workflow content for reconnects. Co-Authored-By: Claude Opus 4.8 (1M context) --- agent-service/src/agent/texera-agent.spec.ts | 109 +++++++++ agent-service/src/agent/texera-agent.ts | 75 ++++++ agent-service/src/server.ts | 59 ++++- agent-service/src/server.ws.spec.ts | 227 ++++++++++++++++++ agent-service/src/types/ws/client.ts | 20 +- agent-service/src/types/ws/server.ts | 23 +- .../agent-chat/agent-chat.component.html | 36 +++ .../agent-chat/agent-chat.component.spec.ts | 99 ++++++++ .../agent-chat/agent-chat.component.ts | 37 +++ .../service/agent/agent.service.spec.ts | 100 ++++++++ .../workspace/service/agent/agent.service.ts | 65 +++++ 11 files changed, 844 insertions(+), 6 deletions(-) create mode 100644 agent-service/src/agent/texera-agent.spec.ts create mode 100644 frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.spec.ts diff --git a/agent-service/src/agent/texera-agent.spec.ts b/agent-service/src/agent/texera-agent.spec.ts new file mode 100644 index 00000000000..3f8b577ec3d --- /dev/null +++ b/agent-service/src/agent/texera-agent.spec.ts @@ -0,0 +1,109 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +// Exercises revert/redo over REAL sendMessage() turns (driven by a mock model), +// so the step-tree, HEAD movement, per-step workflow snapshots, and redo-stack +// lifecycle are covered end to end rather than via hand-built internal state. + +import { describe, expect, test } from "bun:test"; +import { MockLanguageModelV3 } from "ai/test"; +import { TexeraAgent } from "./texera-agent"; + +function textModel(text: string): any { + return new MockLanguageModelV3({ + doGenerate: async () => + ({ + content: [{ type: "text", text }], + finishReason: "stop", + usage: { inputTokens: 1, outputTokens: 1, totalTokens: 2 }, + warnings: [], + }) as any, + }); +} + +function throwingModel(message: string): any { + return new MockLanguageModelV3({ + doGenerate: async () => { + throw new Error(message); + }, + }); +} + +function makeAgent(model: any): TexeraAgent { + return new TexeraAgent({ model, modelType: "mock", agentId: "test-agent" }); +} + +describe("TexeraAgent revert/redo over real turns", () => { + test("a completed turn records a workflow snapshot on every step", async () => { + const agent = makeAgent(textModel("done")); + await agent.sendMessage("hello"); + + const steps = agent.getAllSteps(); + expect(steps.length).toBeGreaterThan(0); + for (const s of steps) { + expect(s.afterWorkflowContent).toBeDefined(); + } + // HEAD is the turn's last step and is on the visible path. + expect(agent.getVisibleReActSteps().at(-1)?.id).toBe(agent.getHead()); + }); + + test("a new prompt clears the redo stack (a branch invalidates redo)", async () => { + const agent = makeAgent(textModel("done")); + await agent.sendMessage("first"); + + const userStep = agent.getAllSteps().find(s => s.role === "user"); + expect(userStep).toBeDefined(); + agent.revertToTurnStart(userStep!.messageId); + expect(agent.canRedo()).toBe(true); + + await agent.sendMessage("second"); // branch from the reverted HEAD + expect(agent.canRedo()).toBe(false); + }); + + test("revert then redo round-trips HEAD across a real turn", async () => { + const agent = makeAgent(textModel("done")); + await agent.sendMessage("only turn"); + const leaf = agent.getHead(); + const userStep = agent.getAllSteps().find(s => s.role === "user")!; + + const reverted = agent.revertToTurnStart(userStep.messageId); + expect(agent.getHead()).toBe(reverted.headId); + expect(agent.getHead()).not.toBe(leaf); + + const redone = agent.redo(); + expect(redone.headId).toBe(leaf); + expect(agent.getHead()).toBe(leaf); + expect(agent.canRedo()).toBe(false); + }); + + test("a turn that errors still snapshots its final step (so it can be reverted/redone)", async () => { + const agent = makeAgent(throwingModel("boom")); + const result = await agent.sendMessage("hello"); + expect(result.error).toBeDefined(); + + const headStep = agent.getStepsById().get(agent.getHead()); + expect(headStep?.afterWorkflowContent).toBeDefined(); + + // The errored turn is still revertable + redoable. + const userStep = agent.getAllSteps().find(s => s.role === "user")!; + agent.revertToTurnStart(userStep.messageId); + const redone = agent.redo(); + expect(redone.workflowContent).toBeDefined(); + }); +}); diff --git a/agent-service/src/agent/texera-agent.ts b/agent-service/src/agent/texera-agent.ts index ccd0545919a..953afee28a9 100644 --- a/agent-service/src/agent/texera-agent.ts +++ b/agent-service/src/agent/texera-agent.ts @@ -91,6 +91,8 @@ export class TexeraAgent { private workflowState: WorkflowState; private metadataStore: WorkflowSystemMetadata; private head: string = INITIAL_STEP_ID; + // Pre-revert HEADs, pushed on revert and cleared on a new prompt. + private redoStack: string[] = []; private stepsById: Map = new Map(); private stepCounter = 0; private workflowResultState: WorkflowResultState; @@ -476,6 +478,7 @@ export class TexeraAgent { } async sendMessage(userMessage: string, messageSource?: "chat" | "feedback"): Promise { + this.redoStack = []; // a new prompt invalidates redo const messageId = `msg-${this.agentId}-${++this.messageCounter}-${Date.now()}`; let stepIndex = 0; @@ -662,6 +665,8 @@ export class TexeraAgent { if (isAborted) { stepIndex++; const stoppedStepId = this.generateStepId(); + // Snapshot so revert/redo work when a turn ends on this step. + const stoppedContent = this.workflowState.getWorkflowContent(); const stoppedStep: ReActStep = { id: stoppedStepId, parentId: this.head, @@ -672,6 +677,8 @@ export class TexeraAgent { content: "Generation stopped by user.", isBegin: false, isEnd: true, + beforeWorkflowContent: stoppedContent, + afterWorkflowContent: stoppedContent, }; this.addStep(stoppedStep); this.head = stoppedStepId; @@ -686,6 +693,8 @@ export class TexeraAgent { stepIndex++; const errorStepId = this.generateStepId(); + // Snapshot so revert/redo work when a turn ends on an error step. + const errorContent = this.workflowState.getWorkflowContent(); const errorStep: ReActStep = { id: errorStepId, parentId: this.head, @@ -696,6 +705,8 @@ export class TexeraAgent { content: `Error: ${error.message || String(error)}`, isBegin: false, isEnd: true, + beforeWorkflowContent: errorContent, + afterWorkflowContent: errorContent, }; this.addStep(errorStep); this.head = errorStepId; @@ -730,10 +741,74 @@ export class TexeraAgent { } } + /** + * Revert the conversation to the state immediately BEFORE the given turn. + * Moves HEAD to that turn's parent step and restores the workflow to the + * turn's pre-edit snapshot. The reverted turn (and any later turns on this + * branch) drop out of the visible path but remain in the tree, so a later + * prompt simply branches from the new HEAD. + * + * @returns the new HEAD id and the restored workflow content + * @throws if the messageId is not a known turn + */ + revertToTurnStart(messageId: string): { headId: string; workflowContent: any } { + const steps = this.reActStepsByMessageId.get(messageId); + if (!steps || steps.length === 0) { + throw new Error(`Unknown turn: ${messageId}`); + } + // steps[0] is the user step: its parentId is the prior HEAD. + const userStep = steps[0]; + const newHead = userStep.parentId ?? INITIAL_STEP_ID; + const workflowContent = userStep.beforeWorkflowContent; + + this.redoStack.push(this.head); + this.head = newHead; + this.currentMessageId = undefined; + if (workflowContent !== undefined) { + this.workflowState.setWorkflowContent(workflowContent); + } + this.log.info({ messageId, newHead }, "reverted to turn start"); + return { headId: newHead, workflowContent }; + } + + /** Whether a previously reverted turn can be redone. */ + canRedo(): boolean { + return this.redoStack.length > 0; + } + + /** + * Redo the most recent revert: move HEAD forward to the step it pointed at + * before that revert, restoring that step's resulting workflow. + * + * @returns the new HEAD id and the restored workflow content + * @throws if there is nothing to redo + */ + redo(): { headId: string; workflowContent: any } { + const target = this.redoStack.pop(); + if (target === undefined) { + throw new Error("Nothing to redo"); + } + this.head = target; + // Fall back to the nearest ancestor snapshot if this step has none. + let workflowContent: any; + let cursor: string | undefined = target; + while (workflowContent === undefined && cursor) { + const step = this.stepsById.get(cursor); + workflowContent = step?.afterWorkflowContent; + cursor = step?.parentId; + } + if (workflowContent !== undefined) { + this.workflowState.setWorkflowContent(workflowContent); + } + this.log.info({ target }, "redid revert"); + return { headId: target, workflowContent }; + } + clearHistory(): void { this.reActStepsByMessageId.clear(); this.stepsById.clear(); this.currentMessageId = undefined; + this.redoStack = []; this.head = INITIAL_STEP_ID; const initialStep: ReActStep = { id: INITIAL_STEP_ID, diff --git a/agent-service/src/server.ts b/agent-service/src/server.ts index 030d27b95bf..394cbcad8e8 100644 --- a/agent-service/src/server.ts +++ b/agent-service/src/server.ts @@ -41,7 +41,13 @@ import type { } from "./types/agent"; import { AgentState, OperatorResultSerializationMode } from "./types/agent"; import type { WsClientCommand, WsServerEvent } from "./types/ws"; -import { WsServerSnapshotEvent, WsServerStepEvent, WsServerStatusEvent, WsServerErrorEvent } from "./types/ws"; +import { + WsServerSnapshotEvent, + WsServerStepEvent, + WsServerStatusEvent, + WsServerErrorEvent, + WsServerHeadChangeEvent, +} from "./types/ws"; import type { OperatorResultSummary } from "./types/execution"; const agentStore = new Map(); @@ -458,7 +464,16 @@ export function buildApp() { agent.addClient(ws); - sendEventToClient(ws, new WsServerSnapshotEvent(agent.getState(), agent.getAllSteps(), agent.getHead())); + sendEventToClient( + ws, + new WsServerSnapshotEvent( + agent.getState(), + agent.getAllSteps(), + agent.getHead(), + agent.getWorkflowState().getWorkflowContent(), + agent.canRedo() + ) + ); }, async message(ws, messageData) { @@ -523,6 +538,42 @@ export function buildApp() { return; } + case "WsClientRevertCommand": { + if (!msg.messageId || typeof msg.messageId !== "string") { + sendEventToClient(ws, new WsServerErrorEvent("messageId is required to revert")); + return; + } + // Don't rewind while a run is in flight or unwinding — it would race + // the loop's own HEAD/workflow mutations (incl. the aborted stopped step). + if (agent.getState() === AgentState.GENERATING || agent.getState() === AgentState.STOPPING) { + sendEventToClient(ws, new WsServerErrorEvent("Cannot revert while the agent is busy")); + return; + } + try { + const { headId, workflowContent } = agent.revertToTurnStart(msg.messageId); + broadcastToAgentClients(agentId, new WsServerHeadChangeEvent(headId, workflowContent, agent.canRedo())); + wsLog.info({ agentId, messageId: msg.messageId, headId }, "reverted turn"); + } catch (error: any) { + sendEventToClient(ws, new WsServerErrorEvent(error.message)); + } + return; + } + + case "WsClientRedoCommand": { + if (agent.getState() === AgentState.GENERATING || agent.getState() === AgentState.STOPPING) { + sendEventToClient(ws, new WsServerErrorEvent("Cannot redo while the agent is busy")); + return; + } + try { + const { headId, workflowContent } = agent.redo(); + broadcastToAgentClients(agentId, new WsServerHeadChangeEvent(headId, workflowContent, agent.canRedo())); + wsLog.info({ agentId, headId }, "redid revert"); + } catch (error: any) { + sendEventToClient(ws, new WsServerErrorEvent(error.message)); + } + return; + } + default: // Frames are parsed from untrusted JSON; reject unknown discriminators // explicitly instead of silently no-op'ing, so client/server mismatches @@ -588,8 +639,10 @@ function printStartupMessage(app: ReturnType) { } console.log(" Send: { type: 'WsClientPromptCommand', content: '...' }"); console.log(" Send: { type: 'WsClientStopCommand' }"); + console.log(" Send: { type: 'WsClientRevertCommand', messageId: '...' }"); + console.log(" Send: { type: 'WsClientRedoCommand' }"); console.log( - " Recv: { type: 'WsServerSnapshotEvent' | 'WsServerStepEvent' | 'WsServerStatusEvent' | 'WsServerErrorEvent', ... }" + " Recv: { type: 'WsServerSnapshotEvent' | 'WsServerStepEvent' | 'WsServerStatusEvent' | 'WsServerErrorEvent' | 'WsServerHeadChangeEvent', ... }" ); } diff --git a/agent-service/src/server.ws.spec.ts b/agent-service/src/server.ws.spec.ts index cbb1ffb8ebd..33b5862a4de 100644 --- a/agent-service/src/server.ws.spec.ts +++ b/agent-service/src/server.ws.spec.ts @@ -159,10 +159,26 @@ describe(`WS ${API}/agents/:id/react`, () => { expect(snapshot.state).toBe("AVAILABLE"); expect(Array.isArray(snapshot.steps)).toBe(true); expect(typeof snapshot.headId).toBe("string"); + // Snapshot carries the current workflow so a (re)connecting client can reload the canvas. + expect(snapshot.workflowContent).toBeDefined(); // Results are pulled on demand, never pushed on the snapshot. expect("operatorResults" in snapshot).toBe(false); }); + test("a client connecting after a revert gets the reverted workflow + canRedo in its snapshot", async () => { + const id = await createAgent(); + const { agent } = seedRevertableTurn(id); + agent.revertToTurnStart("turn-1"); // HEAD -> initial, workflow -> empty, redo available + + const { ws, messages } = connect(id); + await waitOpen(ws); + const snapshot = await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + expect(snapshot.headId).toBe("step-initial"); + expect(snapshot.workflowContent.operators).toHaveLength(0); + expect(snapshot.canRedo).toBe(true); + }); + test("errors and closes when connecting to an unknown agent", async () => { const { messages } = connect("agent-does-not-exist"); const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); @@ -327,4 +343,215 @@ describe(`WS ${API}/agents/:id/react`, () => { expect(ws.readyState).toBe(WebSocket.CLOSED); }); + + // --- Revert (WsClientRevertCommand) --------------------------------------- + + const EMPTY_CONTENT = { operators: [], operatorPositions: {}, links: [], commentBoxes: [], settings: {} }; + const ONE_OP_CONTENT = { + operators: [{ operatorID: "op1" }], + operatorPositions: {}, + links: [], + commentBoxes: [], + settings: {}, + }; + + // Seed a single completed turn ("turn-1") whose pre-edit workflow was empty and + // whose post-edit workflow has one operator; leave HEAD at the turn's last step. + function seedRevertableTurn(id: string) { + const agent = _getAgentForTests(id)! as any; + const priorHead = agent.getHead(); // INITIAL_STEP_ID on a fresh agent + const userStep = { + id: "u1", + parentId: priorHead, + messageId: "turn-1", + stepId: 0, + timestamp: Date.now(), + role: "user", + content: "add an operator", + isBegin: true, + isEnd: true, + beforeWorkflowContent: EMPTY_CONTENT, + afterWorkflowContent: EMPTY_CONTENT, + }; + const agentStep = { + id: "a1", + parentId: "u1", + messageId: "turn-1", + stepId: 1, + timestamp: Date.now(), + role: "agent", + content: "done", + isBegin: true, + isEnd: true, + afterWorkflowContent: ONE_OP_CONTENT, + }; + agent.reActStepsByMessageId.set("turn-1", [userStep, agentStep]); + agent.stepsById.set("u1", userStep); + agent.stepsById.set("a1", agentStep); + agent.head = "a1"; + agent.getWorkflowState().setWorkflowContent(ONE_OP_CONTENT); + return { agent, priorHead }; + } + + test("a revert command rewinds HEAD and broadcasts a head-change frame", async () => { + const id = await createAgent(); + const { agent, priorHead } = seedRevertableTurn(id); + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "turn-1" })); + + const head = await messages.waitFor(m => m.type === "WsServerHeadChangeEvent"); + expect(head.headId).toBe(priorHead); + expect(head.workflowContent.operators).toHaveLength(0); + // The agent truly rewound: HEAD moved and its working workflow is the pre-turn state. + expect(agent.getHead()).toBe(priorHead); + expect(agent.getWorkflowState().getWorkflowContent().operators).toHaveLength(0); + }); + + test("a revert command without a messageId yields an error frame", async () => { + const id = await createAgent(); + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand" })); + + const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(err.error).toBe("messageId is required to revert"); + }); + + test("a revert command for an unknown turn yields an error frame", async () => { + const id = await createAgent(); + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "no-such-turn" })); + + const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(err.error).toBe("Unknown turn: no-such-turn"); + }); + + test("a revert command while generating is rejected", async () => { + const id = await createAgent(); + seedRevertableTurn(id); + const agent = _getAgentForTests(id)! as any; + agent.state = "GENERATING"; + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "turn-1" })); + + const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(err.error).toBe("Cannot revert while the agent is busy"); + // HEAD must not have moved. + expect(agent.getHead()).toBe("a1"); + }); + + test("revert and redo are rejected while STOPPING (aborted run still unwinding)", async () => { + const id = await createAgent(); + const { agent } = seedRevertableTurn(id); + agent.revertToTurnStart("turn-1"); // make redo available + (agent as any).state = "STOPPING"; + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "turn-1" })); + const rerr = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(rerr.error).toBe("Cannot revert while the agent is busy"); + + ws.send(JSON.stringify({ type: "WsClientRedoCommand" })); + const derr = await messages.waitFor(m => m.type === "WsServerErrorEvent" && m.error.includes("redo")); + expect(derr.error).toBe("Cannot redo while the agent is busy"); + }); + + // --- Redo (WsClientRedoCommand) ------------------------------------------- + + test("revert then redo returns HEAD and workflow to the post-turn state", async () => { + const id = await createAgent(); + const { agent } = seedRevertableTurn(id); + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + // Revert: HEAD -> parent, canRedo becomes true. + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "turn-1" })); + const reverted = await messages.waitFor(m => m.type === "WsServerHeadChangeEvent" && m.headId === "step-initial"); + expect(reverted.canRedo).toBe(true); + + // Redo: HEAD -> back to the turn's leaf, workflow restored. + ws.send(JSON.stringify({ type: "WsClientRedoCommand" })); + const redone = await messages.waitFor(m => m.type === "WsServerHeadChangeEvent" && m.headId === "a1"); + expect(redone.workflowContent.operators).toHaveLength(1); + expect(redone.canRedo).toBe(false); + expect(agent.getHead()).toBe("a1"); + expect(agent.getWorkflowState().getWorkflowContent().operators).toHaveLength(1); + }); + + test("redo restores the canvas even when the leaf step has no snapshot (error/stopped turn)", async () => { + const id = await createAgent(); + const { agent } = seedRevertableTurn(id); + // Simulate a turn that ended on an error/stopped step: a leaf with no + // afterWorkflowContent, whose parent (a1) holds the real snapshot. + const errorLeaf = { + id: "err1", + parentId: "a1", + messageId: "turn-1", + stepId: 2, + timestamp: Date.now(), + role: "agent", + content: "Error: rate limit", + isBegin: false, + isEnd: true, + }; + (agent as any).reActStepsByMessageId.get("turn-1").push(errorLeaf); + (agent as any).stepsById.set("err1", errorLeaf); + (agent as any).head = "err1"; + + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRevertCommand", messageId: "turn-1" })); + await messages.waitFor(m => m.type === "WsServerHeadChangeEvent" && m.headId === "step-initial"); + + ws.send(JSON.stringify({ type: "WsClientRedoCommand" })); + const redone = await messages.waitFor(m => m.type === "WsServerHeadChangeEvent" && m.headId === "err1"); + // The leaf carries no snapshot, but redo walks up to a1's snapshot. + expect(redone.workflowContent?.operators).toHaveLength(1); + expect(agent.getWorkflowState().getWorkflowContent().operators).toHaveLength(1); + }); + + test("a redo with nothing to redo yields an error frame", async () => { + const id = await createAgent(); + seedRevertableTurn(id); + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRedoCommand" })); + + const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(err.error).toBe("Nothing to redo"); + }); + + test("a redo while generating is rejected", async () => { + const id = await createAgent(); + const { agent } = seedRevertableTurn(id); + // Make redo available, then flip to GENERATING. + agent.revertToTurnStart("turn-1"); + agent.state = "GENERATING"; + const { ws, messages } = connect(id); + await waitOpen(ws); + await messages.waitFor(m => m.type === "WsServerSnapshotEvent"); + + ws.send(JSON.stringify({ type: "WsClientRedoCommand" })); + + const err = await messages.waitFor(m => m.type === "WsServerErrorEvent"); + expect(err.error).toBe("Cannot redo while the agent is busy"); + }); }); diff --git a/agent-service/src/types/ws/client.ts b/agent-service/src/types/ws/client.ts index 2983ed1d6c5..247c528efd9 100644 --- a/agent-service/src/types/ws/client.ts +++ b/agent-service/src/types/ws/client.ts @@ -36,5 +36,23 @@ export class WsClientStopCommand { readonly type = "WsClientStopCommand"; } +/** + * Revert the workflow to the state BEFORE the given turn and rewind the agent's + * HEAD to that turn's parent step, so a subsequent prompt continues from the + * reverted state. `messageId` identifies the turn (the user-message group). + */ +export class WsClientRevertCommand { + readonly type = "WsClientRevertCommand"; + constructor(readonly messageId: string) {} +} + +/** + * Redo the most recent revert: move HEAD forward to the step it pointed at before + * that revert. Repeatable while the redo stack has entries. Carries no payload. + */ +export class WsClientRedoCommand { + readonly type = "WsClientRedoCommand"; +} + /** Discriminated union of every client -> server frame. */ -export type WsClientCommand = WsClientPromptCommand | WsClientStopCommand; +export type WsClientCommand = WsClientPromptCommand | WsClientStopCommand | WsClientRevertCommand | WsClientRedoCommand; diff --git a/agent-service/src/types/ws/server.ts b/agent-service/src/types/ws/server.ts index 4fdf4e3a06e..5916fabef23 100644 --- a/agent-service/src/types/ws/server.ts +++ b/agent-service/src/types/ws/server.ts @@ -34,7 +34,11 @@ export class WsServerSnapshotEvent { constructor( readonly state: AgentState, readonly steps: ReActStep[], - readonly headId: string + readonly headId: string, + // Workflow at the current HEAD, so a client connecting after a revert/redo can + // reload the canvas to that state instead of only moving its head pointer. + readonly workflowContent: any, + readonly canRedo: boolean = false ) {} } @@ -59,5 +63,20 @@ export class WsServerErrorEvent { constructor(readonly error: string) {} } +/** + * HEAD moved without a new step being produced — emitted after a revert. Carries + * the new HEAD id and the workflow content at that point, so clients can recompute + * the visible step path and reload the canvas. + */ +export class WsServerHeadChangeEvent { + readonly type = "WsServerHeadChangeEvent"; + constructor( + readonly headId: string, + readonly workflowContent: any, + readonly canRedo: boolean = false + ) {} +} + /** Discriminated union of every server -> client frame. */ -export type WsServerEvent = WsServerSnapshotEvent | WsServerStepEvent | WsServerStatusEvent | WsServerErrorEvent; +export type WsServerEvent = + WsServerSnapshotEvent | WsServerStepEvent | WsServerStatusEvent | WsServerErrorEvent | WsServerHeadChangeEvent; diff --git a/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.html b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.html index d650a0a146b..342190339bf 100644 --- a/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.html +++ b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.html @@ -48,6 +48,24 @@ + +
+ +
+
diff --git a/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.spec.ts b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.spec.ts new file mode 100644 index 00000000000..5171adb27f2 --- /dev/null +++ b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.spec.ts @@ -0,0 +1,99 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import { TestBed } from "@angular/core/testing"; +import { AgentChatComponent } from "./agent-chat.component"; +import { AgentService } from "../../../../service/agent/agent.service"; +import { WorkflowActionService } from "../../../../service/workflow-graph/model/workflow-action.service"; +import { NotificationService } from "../../../../../common/service/notification/notification.service"; +import { WorkflowPersistService } from "../../../../../common/service/workflow-persist/workflow-persist.service"; +import { AgentState } from "../../../../service/agent/agent-types"; +import { commonTestProviders } from "../../../../../common/testing/test-utils"; + +describe("AgentChatComponent revert controls", () => { + let component: AgentChatComponent; + const revertTurn = vi.fn(); + const redo = vi.fn(); + + beforeEach(async () => { + revertTurn.mockReset(); + redo.mockReset(); + // compileComponents validates the standalone component's template (including the + // new revert button + nz-popconfirm wiring). The component itself is then + // instantiated directly so the revert-helper logic can be exercised without + // standing up the full render-time DI graph (NzModalService, etc.). + await TestBed.configureTestingModule({ + imports: [AgentChatComponent], + providers: [ + { provide: AgentService, useValue: { revertTurn, redo } }, + { provide: WorkflowActionService, useValue: {} }, + { provide: NotificationService, useValue: { error: () => {}, success: () => {}, warning: () => {} } }, + { provide: WorkflowPersistService, useValue: {} }, + ...commonTestProviders, + ], + }).compileComponents(); + + component = new AgentChatComponent( + { revertTurn, redo } as unknown as AgentService, + {} as WorkflowActionService, + {} as NotificationService, + { detectChanges: () => {} } as any, + {} as WorkflowPersistService + ); + component.agentInfo = { id: "agent-1", name: "Bob", modelType: "gpt-5-mini" } as any; + }); + + it("disallows reverting while generating or stopping, allows it otherwise", () => { + component.agentState = AgentState.GENERATING; + expect(component.canRevert()).toBe(false); + + component.agentState = AgentState.STOPPING; + expect(component.canRevert()).toBe(false); + + component.agentState = AgentState.AVAILABLE; + expect(component.canRevert()).toBe(true); + }); + + it("delegates revertTurn to the agent service with the agent id and turn", () => { + component.revertTurn("msg-7"); + expect(revertTurn).toHaveBeenCalledWith("agent-1", "msg-7"); + }); + + it("allows redo only when redo is available and not generating", () => { + component.canRedoValue = true; + component.agentState = AgentState.AVAILABLE; + expect(component.canRedo()).toBe(true); + + component.agentState = AgentState.GENERATING; + expect(component.canRedo()).toBe(false); + + component.agentState = AgentState.STOPPING; + component.canRedoValue = true; + expect(component.canRedo()).toBe(false); + + component.agentState = AgentState.AVAILABLE; + component.canRedoValue = false; + expect(component.canRedo()).toBe(false); + }); + + it("delegates redo to the agent service", () => { + component.redo(); + expect(redo).toHaveBeenCalledWith("agent-1"); + }); +}); diff --git a/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.ts b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.ts index 55b6c6a3f66..d0feb54ff49 100644 --- a/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.ts +++ b/frontend/src/app/workspace/component/agent/agent-panel/agent-chat/agent-chat.component.ts @@ -59,6 +59,7 @@ import { NzTabsComponent, NzTabComponent } from "ng-zorro-antd/tabs"; import { NzInputNumberComponent } from "ng-zorro-antd/input-number"; import { NzTagComponent } from "ng-zorro-antd/tag"; import { NzSwitchComponent } from "ng-zorro-antd/switch"; +import { NzPopconfirmDirective } from "ng-zorro-antd/popconfirm"; @UntilDestroy() @Component({ @@ -89,6 +90,7 @@ import { NzSwitchComponent } from "ng-zorro-antd/switch"; NzInputGroupComponent, NzInputGroupWhitSuffixOrPrefixDirective, NzSwitchComponent, + NzPopconfirmDirective, ], }) export class AgentChatComponent implements OnInit, AfterViewChecked, OnDestroy, OnChanges { @@ -114,6 +116,8 @@ export class AgentChatComponent implements OnInit, AfterViewChecked, OnDestroy, // Current HEAD step ID in the version tree public currentHeadId: string | null = null; + public canRedoValue = false; + // System info modal state public settingsMaxCharLimit = 20000; // Default max characters for operator results public settingsMaxCellCharLimit = 4000; // Default max characters per cell @@ -208,6 +212,14 @@ export class AgentChatComponent implements OnInit, AfterViewChecked, OnDestroy, this.cdr.detectChanges(); }); + this.agentService + .getCanRedoObservable(this.agentInfo.id) + .pipe(untilDestroyed(this)) + .subscribe(canRedo => { + this.canRedoValue = canRedo; + this.cdr.detectChanges(); + }); + // Subscribe to agent state changes to manage auto-persist // Disable auto-persist when agent is GENERATING, re-enable when AVAILABLE this.agentService @@ -486,6 +498,31 @@ export class AgentChatComponent implements OnInit, AfterViewChecked, OnDestroy, this.agentService.stopGeneration(this.agentInfo.id); } + /** + * Whether a turn can be reverted right now. Reverting mid-generation would race + * the agent's own HEAD/workflow mutations, so it is disabled while generating. + */ + public canRevert(): boolean { + return this.agentState !== AgentState.GENERATING && this.agentState !== AgentState.STOPPING; + } + + /** + * Revert the workflow to the state before the given turn. The agent-service + * rewinds HEAD and pushes a head-change event, which reloads the canvas and + * collapses the chat history to before this turn. + */ + public revertTurn(messageId: string): void { + this.agentService.revertTurn(this.agentInfo.id, messageId); + } + + public canRedo(): boolean { + return this.canRedoValue && this.agentState !== AgentState.GENERATING && this.agentState !== AgentState.STOPPING; + } + + public redo(): void { + this.agentService.redo(this.agentInfo.id); + } + public clearMessages(): void { this.agentService.clearMessages(this.agentInfo.id); } diff --git a/frontend/src/app/workspace/service/agent/agent.service.spec.ts b/frontend/src/app/workspace/service/agent/agent.service.spec.ts index 1f9bcd82591..8912cd4bb6b 100644 --- a/frontend/src/app/workspace/service/agent/agent.service.spec.ts +++ b/frontend/src/app/workspace/service/agent/agent.service.spec.ts @@ -149,4 +149,104 @@ describe("AgentService", () => { .flush({ status: "stopping" }); }); }); + + describe("revertTurn", () => { + it("sends a revert command for the given turn over the websocket", () => { + const send = vi.fn(); + (service as any).agentStateTracking.set("agent-1", { + websocket: { readyState: WebSocket.OPEN, send }, + }); + + service.revertTurn("agent-1", "msg-7"); + + expect(send).toHaveBeenCalledWith(JSON.stringify({ type: "WsClientRevertCommand", messageId: "msg-7" })); + }); + + it("does not send when no websocket is open", () => { + const send = vi.fn(); + (service as any).agentStateTracking.set("agent-1", { + websocket: { readyState: WebSocket.CLOSED, send }, + }); + + service.revertTurn("agent-1", "msg-7"); + + expect(send).not.toHaveBeenCalled(); + }); + }); + + describe("redo", () => { + it("sends a redo command over the websocket", () => { + const send = vi.fn(); + (service as any).agentStateTracking.set("agent-1", { + websocket: { readyState: WebSocket.OPEN, send }, + }); + + service.redo("agent-1"); + + expect(send).toHaveBeenCalledWith(JSON.stringify({ type: "WsClientRedoCommand" })); + }); + }); + + describe("sendMessage clears redo", () => { + it("resets canRedo when a new prompt is sent (mirrors the backend clearing its stack)", () => { + const send = vi.fn(); + const tracking = (service as any).getOrCreateStateTracking("agent-1"); + tracking.websocket = { readyState: WebSocket.OPEN, send }; + tracking.canRedoSubject.next(true); + (service as any).agents.set("agent-1", { id: "agent-1", name: "Bob" }); + + service.sendMessage("agent-1", "hello"); + + expect(send).toHaveBeenCalled(); + expect(tracking.canRedoSubject.getValue()).toBe(false); + }); + }); + + describe("WsServerHeadChangeEvent handling", () => { + it("advances the head pointer and reloads the workflow on a revert", () => { + const tracking = (service as any).getOrCreateStateTracking("agent-1"); + const content = { operators: [], operatorPositions: {}, links: [], commentBoxes: [], settings: {} }; + + (service as any).handleWebSocketMessage("agent-1", tracking, { + type: "WsServerHeadChangeEvent", + headId: "step-initial", + workflowContent: content, + }); + + expect(tracking.headIdSubject.getValue()).toEqual("step-initial"); + expect(tracking.workflowSubject.getValue()?.content).toEqual(content); + expect(tracking.wsWorkflowActive).toBe(true); + }); + + it("applies workflowContent and canRedo from a snapshot (reconnect after revert)", () => { + const tracking = (service as any).getOrCreateStateTracking("agent-1"); + const content = { operators: [], operatorPositions: {}, links: [], commentBoxes: [], settings: {} }; + + (service as any).handleWebSocketMessage("agent-1", tracking, { + type: "WsServerSnapshotEvent", + state: "AVAILABLE", + steps: [], + headId: "step-initial", + workflowContent: content, + canRedo: true, + }); + + expect(tracking.headIdSubject.getValue()).toBe("step-initial"); + expect(tracking.workflowSubject.getValue()?.content).toEqual(content); + expect(tracking.canRedoSubject.getValue()).toBe(true); + }); + + it("updates canRedo from the head-change event", () => { + const tracking = (service as any).getOrCreateStateTracking("agent-1"); + + (service as any).handleWebSocketMessage("agent-1", tracking, { + type: "WsServerHeadChangeEvent", + headId: "step-initial", + workflowContent: { operators: [], operatorPositions: {}, links: [], commentBoxes: [], settings: {} }, + canRedo: true, + }); + + expect(tracking.canRedoSubject.getValue()).toBe(true); + }); + }); }); diff --git a/frontend/src/app/workspace/service/agent/agent.service.ts b/frontend/src/app/workspace/service/agent/agent.service.ts index 5e7c254f22c..63d9c1e5952 100644 --- a/frontend/src/app/workspace/service/agent/agent.service.ts +++ b/frontend/src/app/workspace/service/agent/agent.service.ts @@ -172,6 +172,8 @@ interface AgentStateTracking { }>; /** Current HEAD step ID in the version tree */ headIdSubject: BehaviorSubject; + /** Whether a reverted turn can be redone (server-tracked redo stack) */ + canRedoSubject: BehaviorSubject; workflowSubject: BehaviorSubject; workflowId?: number; stopPolling$: Subject; @@ -364,6 +366,7 @@ export class AgentService { modifiedOperatorIds: string[]; }>({ viewedOperatorIds: [], addedOperatorIds: [], modifiedOperatorIds: [] }), headIdSubject: new BehaviorSubject(null), + canRedoSubject: new BehaviorSubject(false), workflowSubject: new BehaviorSubject(null), workflowId, stopPolling$: new Subject(), @@ -462,6 +465,7 @@ export class AgentService { if (message.headId !== undefined) { tracking.headIdSubject.next(message.headId); } + tracking.canRedoSubject.next(message.canRedo === true); // Handle initial workflow content from agent service (ground truth) if (message.workflowContent) { tracking.wsWorkflowActive = true; @@ -535,6 +539,23 @@ export class AgentService { } break; + case "WsServerHeadChangeEvent": + // Revert/redo: advancing head recomputes visible steps; the workflow reloads the canvas. + if (message.headId !== undefined) { + tracking.headIdSubject.next(message.headId); + } + tracking.canRedoSubject.next(message.canRedo === true); + if (message.workflowContent !== undefined) { + tracking.wsWorkflowActive = true; + const existingWorkflow = tracking.workflowSubject.getValue(); + const workflow = { + ...(existingWorkflow || {}), + content: message.workflowContent, + } as Workflow; + tracking.workflowSubject.next(workflow); + } + break; + default: console.warn("Unknown agent WebSocket message type:", message.type); } @@ -877,6 +898,8 @@ export class AgentService { try { tracking.websocket.send(JSON.stringify(wsMessage)); + // Only mirror the backend's redo-stack clear once the prompt actually went out. + tracking.canRedoSubject.next(false); } catch (error) { console.error("Failed to send message to agent:", error); this.notificationService.error("Failed to send message"); @@ -920,6 +943,48 @@ export class AgentService { }); } + /** + * Revert the workflow to the state before the given turn. The agent-service + * rewinds its HEAD and replies with a WsServerHeadChangeEvent, which updates + * the head pointer (recomputing the visible steps) and reloads the canvas. + */ + public revertTurn(agentId: string, messageId: string): void { + const tracking = this.agentStateTracking.get(agentId); + if (!tracking || !tracking.websocket || tracking.websocket.readyState !== WebSocket.OPEN) { + this.notificationService.error("WebSocket connection not available"); + return; + } + try { + tracking.websocket.send(JSON.stringify({ type: "WsClientRevertCommand", messageId })); + } catch (error) { + console.error("Failed to send revert command:", error); + this.notificationService.error("Failed to revert"); + } + } + + /** + * Redo the most recent revert. The agent-service moves HEAD forward and replies + * with a WsServerHeadChangeEvent, which reloads the canvas and updates canRedo. + */ + public redo(agentId: string): void { + const tracking = this.agentStateTracking.get(agentId); + if (!tracking || !tracking.websocket || tracking.websocket.readyState !== WebSocket.OPEN) { + this.notificationService.error("WebSocket connection not available"); + return; + } + try { + tracking.websocket.send(JSON.stringify({ type: "WsClientRedoCommand" })); + } catch (error) { + console.error("Failed to send redo command:", error); + this.notificationService.error("Failed to redo"); + } + } + + /** Observable of whether a reverted turn can currently be redone. */ + public getCanRedoObservable(agentId: string): Observable { + return this.getOrCreateStateTracking(agentId).canRedoSubject.asObservable(); + } + /** * Stop generation for an agent via WebSocket. */