feat(agent): add workflow revert & redo for AI agent turns

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();
+  });
+});

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<string, ReActStep> = new Map();
   private stepCounter = 0;
   private workflowResultState: WorkflowResultState;
@@ -476,6 +478,7 @@ export class TexeraAgent {
   }
 
   async sendMessage(userMessage: string, messageSource?: "chat" | "feedback"): Promise<AgentMessageResult> {
+    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,

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<string, TexeraAgent>();
@@ -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<typeof buildApp>) {
     }
     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', ... }"
     );
   }