feat(agent-service): persist agent state to disk and rehydrate on startup

agent-service/.env.example

@@ -18,4 +18,10 @@ LLM_ENDPOINT=http://localhost:9096
 # Texera backend services
 TEXERA_DASHBOARD_SERVICE_ENDPOINT=http://localhost:8080
 WORKFLOW_COMPILING_SERVICE_ENDPOINT=http://localhost:9090
-WORKFLOW_EXECUTION_SERVICE_ENDPOINT=http://localhost:8085
+WORKFLOW_EXECUTION_SERVICE_ENDPOINT=http://localhost:8085
+
+# Directory where agent state (conversation, workflow, settings) is persisted as
+# one JSON file per agent. Agents are rehydrated from here on startup, so they
+# survive restarts. Leave empty to disable persistence (agents live only in
+# memory, the previous behavior). Mount a volume here in containerized deploys.
+AGENT_STATE_DIR=

agent-service/src/agent/texera-agent-snapshot.test.ts

@@ -0,0 +1,169 @@
+/**
+ * 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 { describe, expect, test } from "bun:test";
+import { TexeraAgent } from "./texera-agent";
+import type { LanguageModel } from "ai";
+import { INITIAL_STEP_ID, OperatorResultSerializationMode } from "../types/agent";
+import type { AgentSnapshot } from "../types/agent";
+import type { WorkflowContent } from "../types/workflow";
+
+function newAgent(createdAt?: Date): TexeraAgent {
+  return new TexeraAgent({
+    model: {} as LanguageModel,
+    modelType: "test-model",
+    agentId: "agent-x",
+    agentName: "Bob",
+    createdAt,
+  });
+}
+
+const sampleWorkflow: WorkflowContent = {
+  operators: [
+    {
+      operatorID: "op1",
+      operatorType: "Filter",
+      operatorVersion: "1.0",
+      operatorProperties: { predicate: "x > 1" },
+      inputPorts: [{ portID: "input-0", displayName: "in" }],
+      outputPorts: [{ portID: "output-0", displayName: "out" }],
+      showAdvanced: false,
+    },
+  ],
+  operatorPositions: { op1: { x: 10, y: 20 } },
+  links: [],
+  commentBoxes: [],
+  settings: { dataTransferBatchSize: 400 },
+};
+
+function sampleSnapshot(): AgentSnapshot {
+  return {
+    version: 1,
+    agentId: "agent-x",
+    agentName: "Bob",
+    modelType: "test-model",
+    createdAt: "2024-01-01T00:00:00.000Z",
+    head: "s1",
+    stepCounter: 3,
+    messageCounter: 2,
+    settings: {
+      disabledTools: ["executeOperator"],
+      maxOperatorResultCharLimit: 1500,
+      maxOperatorResultCellCharLimit: 800,
+      operatorResultSerializationMode: OperatorResultSerializationMode.TSV,
+      toolTimeoutMs: 120000,
+      executionTimeoutMs: 180000,
+      maxSteps: 42,
+      allowedOperatorTypes: ["Filter", "CSVFileScan"],
+    },
+    delegate: {
+      userInfo: { uid: 1, name: "user", email: "[email protected]", role: "REGULAR" },
+      workflowId: 7,
+      workflowName: "wf",
+      computingUnitId: 3,
+    },
+    steps: [
+      {
+        id: INITIAL_STEP_ID,
+        messageId: "initial",
+        stepId: -1,
+        timestamp: 0,
+        role: "user",
+        content: "",
+        isBegin: true,
+        isEnd: true,
+      },
+      {
+        id: "s1",
+        parentId: INITIAL_STEP_ID,
+        messageId: "m1",
+        stepId: 0,
+        timestamp: 123,
+        role: "user",
+        content: "build a filter",
+        isBegin: true,
+        isEnd: true,
+      },
+    ],
+    messageGroups: { m1: ["s1"] },
+    workflowContent: sampleWorkflow,
+  };
+}
+
+describe("TexeraAgent.toSnapshot", () => {
+  test("captures a fresh agent: version, empty history, initial head", () => {
+    const snap = newAgent().toSnapshot();
+    expect(snap.version).toBe(1);
+    expect(snap.agentId).toBe("agent-x");
+    expect(snap.head).toBe(INITIAL_STEP_ID);
+    // Only the sentinel initial step exists.
+    expect(snap.steps).toHaveLength(1);
+    expect(snap.steps[0].id).toBe(INITIAL_STEP_ID);
+    expect(snap.messageGroups).toEqual({});
+    expect(snap.delegate).toBeUndefined();
+  });
+
+  test("preserves createdAt as an ISO string", () => {
+    const snap = newAgent(new Date("2024-01-01T00:00:00.000Z")).toSnapshot();
+    expect(snap.createdAt).toBe("2024-01-01T00:00:00.000Z");
+  });
+});
+
+describe("TexeraAgent.restoreFromSnapshot", () => {
+  test("restores conversation, settings, workflow, and delegate", () => {
+    const snap = sampleSnapshot();
+    const agent = newAgent(new Date(snap.createdAt));
+    agent.restoreFromSnapshot(snap);
+
+    expect(agent.getHead()).toBe("s1");
+    expect(agent.getReActSteps().map(s => s.id)).toEqual(["s1"]);
+    expect(agent.getAllSteps().map(s => s.id)).toEqual(["s1"]); // excludes the initial sentinel
+
+    const settings = agent.getSettings();
+    expect(settings.maxSteps).toBe(42);
+    expect(settings.maxOperatorResultCharLimit).toBe(1500);
+    expect(Array.from(settings.disabledTools)).toEqual(["executeOperator"]);
+
+    expect(
+      agent
+        .getWorkflowState()
+        .getWorkflowContent()
+        .operators.map(o => o.operatorID)
+    ).toEqual(["op1"]);
+
+    // Delegate metadata is restored, but the user token is not persisted.
+    expect(agent.getDelegateConfig()?.workflowId).toBe(7);
+    expect(agent.getDelegateConfig()?.userToken).toBe("");
+  });
+
+  test("round-trips: toSnapshot(restore(s)) deep-equals s", () => {
+    const snap = sampleSnapshot();
+    const agent = newAgent(new Date(snap.createdAt));
+    agent.restoreFromSnapshot(snap);
+
+    // Survives a JSON round-trip as it would on disk.
+    const roundTripped = JSON.parse(JSON.stringify(agent.toSnapshot())) as AgentSnapshot;
+    expect(roundTripped).toEqual(snap);
+  });
+
+  test("rejects an unsupported snapshot version", () => {
+    const snap = { ...sampleSnapshot(), version: 2 as unknown as 1 };
+    expect(() => newAgent().restoreFromSnapshot(snap)).toThrow(/Unsupported agent snapshot version/);
+  });
+});

agent-service/src/agent/texera-agent.ts

@@ -24,7 +24,7 @@ import { WorkflowState } from "./workflow-state";
 import { WorkflowSystemMetadata } from "./util/workflow-system-metadata";
 import { WorkflowResultState } from "./workflow-result-state";
 import { formatOperatorResult } from "./tools/result-formatting";
-import type { AgentSettings, ReActStep, TokenUsage, UserInfo } from "../types/agent";
+import type { AgentSettings, AgentSnapshot, ReActStep, TokenUsage, UserInfo } from "../types/agent";
 import {
   AgentState as AgentStateEnum,
   DEFAULT_AGENT_SETTINGS,
@@ -60,6 +60,8 @@ export interface TexeraAgentConfig {
   agentId: string;
   agentName?: string;
   systemPrompt?: string;
+  // Preserved across restarts when an agent is reconstructed from a snapshot.
+  createdAt?: Date;
 }
 
 export interface AgentMessageResult {
@@ -129,7 +131,7 @@ export class TexeraAgent {
     this.agentId = config.agentId;
     this.agentName = config.agentName || `Agent-${config.agentId}`;
     this.modelType = config.modelType;
-    this.createdAt = new Date();
+    this.createdAt = config.createdAt ?? new Date();
     this.model = config.model;
     this.systemPrompt = config.systemPrompt || "";
     this.log = createLogger("TexeraAgent", { agentId: this.agentId });
@@ -823,6 +825,107 @@ export class TexeraAgent {
     return relevantSteps;
   }
 
+  /**
+   * Produce a durable, JSON-serializable snapshot of this agent. The user token
+   * and execution-result caches are intentionally excluded (see AgentSnapshot).
+   */
+  toSnapshot(): AgentSnapshot {
+    const messageGroups: Record<string, string[]> = {};
+    for (const [messageId, steps] of this.reActStepsByMessageId) {
+      messageGroups[messageId] = steps.map(s => s.id);
+    }
+
+    return {
+      version: 1,
+      agentId: this.agentId,
+      agentName: this.agentName,
+      modelType: this.modelType,
+      createdAt: this.createdAt.toISOString(),
+      head: this.head,
+      stepCounter: this.stepCounter,
+      messageCounter: this.messageCounter,
+      settings: {
+        disabledTools: Array.from(this.settings.disabledTools),
+        maxOperatorResultCharLimit: this.settings.maxOperatorResultCharLimit,
+        maxOperatorResultCellCharLimit: this.settings.maxOperatorResultCellCharLimit,
+        operatorResultSerializationMode: this.settings.operatorResultSerializationMode,
+        toolTimeoutMs: this.settings.toolTimeoutMs,
+        executionTimeoutMs: this.settings.executionTimeoutMs,
+        maxSteps: this.settings.maxSteps,
+        allowedOperatorTypes: this.settings.allowedOperatorTypes,
+      },
+      delegate: this.delegateConfig
+        ? {
+            userInfo: this.delegateConfig.userInfo,
+            workflowId: this.delegateConfig.workflowId,
+            workflowName: this.delegateConfig.workflowName,
+            computingUnitId: this.delegateConfig.computingUnitId,
+          }
+        : undefined,
+      steps: Array.from(this.stepsById.values()),
+      messageGroups,
+      workflowContent: this.workflowState.getWorkflowContent(),
+    };
+  }
+
+  /**
+   * Restore conversation, workflow, settings, and delegate metadata from a
+   * snapshot. Must be called on a freshly constructed agent (one built with the
+   * snapshot's createdAt). The restored delegate carries no user token, so the
+   * caller must re-attach one before the agent can execute or auto-persist.
+   */
+  restoreFromSnapshot(snapshot: AgentSnapshot): void {
+    if (snapshot.version !== 1) {
+      throw new Error(`Unsupported agent snapshot version: ${snapshot.version}`);
+    }
+
+    this.stepCounter = snapshot.stepCounter;
+    this.messageCounter = snapshot.messageCounter;
+
+    this.settings = {
+      ...DEFAULT_AGENT_SETTINGS,
+      ...snapshot.settings,
+      disabledTools: new Set(snapshot.settings.disabledTools),
+      systemPrompt: this.systemPrompt,
+    };
+
+    this.stepsById = new Map(snapshot.steps.map(step => [step.id, step]));
+    this.reActStepsByMessageId = new Map();
+    for (const [messageId, stepIds] of Object.entries(snapshot.messageGroups)) {
+      const steps = stepIds.map(id => this.stepsById.get(id)).filter((s): s is ReActStep => s !== undefined);
+      this.reActStepsByMessageId.set(messageId, steps);
+    }
+
+    // Ensure the initial sentinel step always exists so HEAD traversal/checkout works.
+    if (!this.stepsById.has(INITIAL_STEP_ID)) {
+      this.stepsById.set(INITIAL_STEP_ID, {
+        id: INITIAL_STEP_ID,
+        messageId: "initial",
+        stepId: -1,
+        timestamp: Date.now(),
+        role: "user",
+        content: "",
+        isBegin: true,
+        isEnd: true,
+      });
+    }
+    this.head = snapshot.head;
+
+    if (snapshot.delegate) {
+      this.delegateConfig = {
+        userToken: "",
+        userInfo: snapshot.delegate.userInfo,
+        workflowId: snapshot.delegate.workflowId,
+        workflowName: snapshot.delegate.workflowName,
+        computingUnitId: snapshot.delegate.computingUnitId,
+      };
+    }
+
+    this.workflowState.setWorkflowContent(snapshot.workflowContent);
+    this.rebuildSystemPrompt();
+    this.tools = this.createTools();
+  }
+
   destroy(): void {
     if (this.workflowChangeSubscription) {
       this.workflowChangeSubscription.unsubscribe();