refactor(agent-service): reconcile rebased stack with origin/main

Reconciliation for rebasing the agent-service-execution-model stack onto
origin/main, which had independently merged #5751 (extract WS types into
types/ws as classes) and #5930 (remove the dead checkout/headChange path).
Our stack redesigns the same WS layer as interface discriminated unions, so
the two diverge structurally. Changes:

- Restore the execution-model redesign in types/execution.ts: the 3-way merge
  re-appended main's #5751 flat `OperatorResultSummary` (which references the
  removed `PortShape` and duplicates our redesigned type), so drop it; our
  redesign's `OperatorResultSummary`/`OperatorExecutionSummary` supersede it.
- Adopt #5751 in server.ts: the TexeraAgent websocket surface was renamed
  `websockets` -> `clients` (getClients/addClient/removeClient) by the
  auto-merge in texera-agent.ts; align the server call sites, restore the
  `_getAgentForTests` test accessor, and reject unknown ws message types with
  an error frame (matching main's handler) instead of silently ignoring them.
- Adopt #5930: drop the checkout/headChange dead path that our stack still
  carried — the POST /:id/checkout route + headChange broadcast and the
  startup-banner token in server.ts, `WsServerHeadChangeMessage` (and its now
  orphaned WorkflowContent import) in types/ws/server.ts, and the
  `case "headChange"` handler + `checkoutStep()` in the frontend agent.service.
  (TexeraAgent.checkout() was already removed by the auto-merge.)
- Reconcile tests to the new discriminated-union ws protocol: rewrite
  server.ws.spec.ts (init/step/state/complete/error frames; prompt/command
  client frames; operatorResults streamed on init/complete), update the
  server.spec.ts operator-results test to the OperatorExecutionSummary shape,
  and fix the frontend stopGeneration assertion to the {type:"command",
  commandType:"stop"} wire format.

agent-service: tsc --noEmit clean, 143/143 bun test pass, prettier clean.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: bobbai00

agent-service/src/server.spec.ts

@@ -323,34 +323,33 @@ describe("agent read routes", () => {
 
   test("GET /:id/operator-results maps the visible operator results", async () => {
     const agent = _getAgentForTests(id)!;
+    // The route returns each operator's OperatorExecutionSummary verbatim.
     (agent as any).getWorkflowResultState = () => ({
       getAllVisible: () =>
         new Map([
           [
             "op-1",
             {
               operatorInfo: {
-                state: "COMPLETED",
-                inputTuples: 1,
-                outputTuples: 2,
-                inputPortShapes: [],
-                result: [{ a: 1 }],
-                error: undefined,
-                warnings: [],
-                consoleLogs: [],
-                totalRowCount: 2,
-                resultStatistics: {},
+                state: "Completed",
+                errorMessages: [],
+                resultSummary: {
+                  resultMode: "table",
+                  sampleTuples: [{ a: 1 }],
+                  totalRowCount: 2,
+                },
               },
             },
           ],
         ]),
     });
 
-    const body = await readJson<{ results: Record<string, { outputTuples: number; outputColumns: number }> }>(
-      await getJson(`${API}/agents/${id}/operator-results`)
-    );
-    expect(body.results["op-1"].outputTuples).toBe(2);
-    expect(body.results["op-1"].outputColumns).toBe(1);
+    const body = await readJson<{
+      results: Record<string, { state: string; resultSummary: { totalRowCount: number; sampleTuples: unknown[] } }>;
+    }>(await getJson(`${API}/agents/${id}/operator-results`));
+    expect(body.results["op-1"].state).toBe("Completed");
+    expect(body.results["op-1"].resultSummary.totalRowCount).toBe(2);
+    expect(body.results["op-1"].resultSummary.sampleTuples).toEqual([{ a: 1 }]);
   });
 });
 

agent-service/src/server.ts

@@ -309,31 +309,6 @@ const agentsRouter = new Elysia({ prefix: "/agents" })
     return { status: "cleared" };
   })
 
-  .post("/:id/checkout", ({ params: { id }, body }) => {
-    const agent = getAgent(id);
-    const { stepId } = body as { stepId: string };
-    if (!stepId) throw new Error("stepId is required");
-
-    const success = agent.checkout(stepId);
-    if (!success) throw new Error(`Step ${stepId} not found or checkout failed`);
-
-    const allSteps = agent.getAllSteps();
-    const workflowContent = agent.getWorkflowState().getWorkflowContent();
-
-    broadcastToAgent(id, {
-      type: "headChange",
-      headId: stepId,
-      steps: allSteps,
-      workflowContent,
-      operatorResults: getOperatorResultSummaries(agent),
-    });
-
-    return {
-      status: "checked out",
-      headId: stepId,
-    };
-  })
-
   .get("/:id/operator-types", ({ params: { id } }) => {
     const agent = getAgent(id);
     const metadataStore = agent.getMetadataStore();
@@ -426,12 +401,12 @@ function broadcastToAgent(agentId: string, message: WsServerMessage): void {
   if (!agent) return;
 
   const jsonMessage = JSON.stringify(message);
-  for (const ws of agent.getWebsockets()) {
+  for (const ws of agent.getClients()) {
     try {
       ws.send(jsonMessage);
     } catch (error) {
       wsLog.error({ agentId, err: error }, "failed to send message to client");
-      agent.removeWebsocket(ws);
+      agent.removeClient(ws);
     }
   }
 }
@@ -459,7 +434,7 @@ export function buildApp() {
           return;
         }
 
-        agent.addWebsocket(ws);
+        agent.addClient(ws);
 
         const initMessage: WsServerInitMessage = {
           type: "init",
@@ -488,6 +463,12 @@ export function buildApp() {
           return;
         }
 
+        if (msg.type !== "command" && msg.type !== "prompt") {
+          const unknownType = (msg as { type?: string }).type;
+          ws.send(JSON.stringify({ type: "error", error: `Unknown message type: ${unknownType}` }));
+          return;
+        }
+
         if (msg.type === "command") {
           if (msg.commandType === "stop") {
             agent.stop();
@@ -546,7 +527,7 @@ export function buildApp() {
 
         const agent = agentStore.get(agentId);
         if (agent) {
-          agent.removeWebsocket(ws);
+          agent.removeClient(ws);
         }
       },
     })
@@ -564,6 +545,10 @@ export function _resetAgentStoreForTests(): void {
   agentCounter = 0;
 }
 
+export function _getAgentForTests(agentId: string): TexeraAgent | undefined {
+  return agentStore.get(agentId);
+}
+
 function printStartupMessage(app: ReturnType<typeof buildApp>) {
   const LINE = "=".repeat(60);
   console.log(LINE);
@@ -591,7 +576,7 @@ function printStartupMessage(app: ReturnType<typeof buildApp>) {
     }
     console.log("         Send: { type: 'prompt', content: '...' }");
     console.log("         Send: { type: 'command', commandType: 'stop' }");
-    console.log("         Recv: { type: 'step' | 'state' | 'complete' | 'error' | 'init' | 'headChange', ... }");
+    console.log("         Recv: { type: 'step' | 'state' | 'complete' | 'error' | 'init', ... }");
   }
 
   console.log("");

agent-service/src/server.ws.spec.ts

@@ -17,11 +17,13 @@
  * under the License.
  */
 
-// Exercises the /agents/:id/react WebSocket protocol end to end: the snapshot
-// sent on connect, the status lifecycle frames, the stop command, the prompt
-// request (with a stubbed run), and the error paths. These drive the real
-// socket via app.listen + a WebSocket client, since app.handle() does not
-// perform WS upgrades.
+// Exercises the /agents/:id/react WebSocket protocol end to end: the init
+// snapshot sent on connect, the state lifecycle frames, the stop command, the
+// prompt request (with a stubbed run), and the error paths. Frames use the
+// discriminated-union message model (type: "init" | "step" | "state" |
+// "complete" | "error" for server frames; "prompt" | "command" for client
+// frames). These drive the real socket via app.listen + a WebSocket client,
+// since app.handle() does not perform WS upgrades.
 
 import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, test } from "bun:test";
 import { buildApp, _resetAgentStoreForTests, _getAgentForTests } from "./server";
@@ -150,74 +152,75 @@ afterEach(() => {
 });
 
 describe(`WS ${API}/agents/:id/react`, () => {
-  test("sends a results-free snapshot frame on connect", async () => {
+  test("sends an init snapshot frame on connect", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
 
-    const snapshot = await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    const snapshot = await messages.waitFor(m => m.type === "init");
     expect(snapshot.state).toBe("AVAILABLE");
     expect(Array.isArray(snapshot.steps)).toBe(true);
     expect(typeof snapshot.headId).toBe("string");
-    // Results are pulled on demand, never pushed on the snapshot.
-    expect("operatorResults" in snapshot).toBe(false);
+    // The init frame carries the current per-operator execution summaries.
+    expect("operatorResults" in snapshot).toBe(true);
+    expect(typeof snapshot.operatorResults).toBe("object");
   });
 
   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");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("Agent not found");
   });
 
-  test("a stop command broadcasts a STOPPING status frame", async () => {
+  test("a stop command broadcasts a STOPPING state frame", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
-    ws.send(JSON.stringify({ type: "WsClientStopCommand" }));
+    ws.send(JSON.stringify({ type: "command", commandType: "stop" }));
 
-    const status = await messages.waitFor(m => m.type === "WsServerStatusEvent");
+    const status = await messages.waitFor(m => m.type === "state");
     expect(status.state).toBe("STOPPING");
   });
 
   test("a prompt with empty content yields an error frame", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
-    ws.send(JSON.stringify({ type: "WsClientPromptCommand", content: "" }));
+    ws.send(JSON.stringify({ type: "prompt", content: "" }));
 
-    const err = await messages.waitFor(m => m.type === "WsServerErrorEvent");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("Message content is required");
   });
 
   test("a malformed (non-JSON) frame yields an error frame", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
     ws.send("this is not json");
 
-    const err = await messages.waitFor(m => m.type === "WsServerErrorEvent");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("Invalid message format");
   });
 
   test("an unknown message type yields an error frame", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
     ws.send(JSON.stringify({ type: "bogus" }));
 
-    const err = await messages.waitFor(m => m.type === "WsServerErrorEvent");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("Unknown message type: bogus");
   });
 
-  test("a prompt run streams GENERATING -> step -> resting status (no result frames)", async () => {
+  test("a prompt run streams GENERATING -> step -> complete", async () => {
     const id = await createAgent();
 
     // Stub the agent's run so no live LLM is needed: emit one ending step via
@@ -259,22 +262,25 @@ describe(`WS ${API}/agents/:id/react`, () => {
 
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
-    ws.send(JSON.stringify({ type: "WsClientPromptCommand", content: "hello" }));
+    ws.send(JSON.stringify({ type: "prompt", content: "hello" }));
 
-    const generating = await messages.waitFor(m => m.type === "WsServerStatusEvent" && m.state === "GENERATING");
+    const generating = await messages.waitFor(m => m.type === "state" && m.state === "GENERATING");
     expect(generating.state).toBe("GENERATING");
 
-    const step = await messages.waitFor(m => m.type === "WsServerStepEvent");
+    const step = await messages.waitFor(m => m.type === "step");
     expect(step.step.content).toBe("done");
+    // A step with no tool calls carries no operatorResults.
     expect("operatorResults" in step).toBe(false);
 
-    const resting = await messages.waitFor(m => m.type === "WsServerStatusEvent" && m.state === "AVAILABLE");
-    expect(resting.state).toBe("AVAILABLE");
+    // The run ends with a terminal `complete` frame carrying the resting state.
+    const complete = await messages.waitFor(m => m.type === "complete");
+    expect(complete.state).toBe("AVAILABLE");
+    expect("operatorResults" in complete).toBe(true);
   });
 
-  test("a failed run emits an error frame and still returns to a resting status", async () => {
+  test("a failed run emits an error frame", async () => {
     const id = await createAgent();
 
     const agent = _getAgentForTests(id)!;
@@ -284,40 +290,35 @@ describe(`WS ${API}/agents/:id/react`, () => {
 
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
-    ws.send(JSON.stringify({ type: "WsClientPromptCommand", content: "hello" }));
+    ws.send(JSON.stringify({ type: "prompt", content: "hello" }));
 
-    await messages.waitFor(m => m.type === "WsServerStatusEvent" && m.state === "GENERATING");
+    await messages.waitFor(m => m.type === "state" && m.state === "GENERATING");
 
-    const err = await messages.waitFor(m => m.type === "WsServerErrorEvent");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("boom");
-
-    // The end-of-run status frame must still fire after a failure, so the client
-    // is not left stuck on GENERATING.
-    const resting = await messages.waitFor(m => m.type === "WsServerStatusEvent" && m.state === "AVAILABLE");
-    expect(resting.state).toBe("AVAILABLE");
   });
 
   test("a message for an agent that no longer exists yields an error frame", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
     // Drop the agent while the socket stays open; the message handler re-looks it up.
     _resetAgentStoreForTests();
-    ws.send(JSON.stringify({ type: "WsClientPromptCommand", content: "hello" }));
+    ws.send(JSON.stringify({ type: "prompt", content: "hello" }));
 
-    const err = await messages.waitFor(m => m.type === "WsServerErrorEvent");
+    const err = await messages.waitFor(m => m.type === "error");
     expect(err.error).toBe("Agent not found");
   });
 
   test("runs the close handler when the client disconnects", async () => {
     const id = await createAgent();
     const { ws, messages } = connect(id);
     await waitOpen(ws);
-    await messages.waitFor(m => m.type === "WsServerSnapshotEvent");
+    await messages.waitFor(m => m.type === "init");
 
     const closed = new Promise<void>(resolve => ws.addEventListener("close", () => resolve(), { once: true }));
     ws.close();

agent-service/src/types/execution.ts

@@ -95,22 +95,3 @@ export interface WorkflowExecutionSummary {
   // empty means none.
   errors: string[];
 }
-
-/**
- * Wire projection of one operator's execution result, summarized for the
- * client: counts and a small record sample instead of full payloads. Returned
- * by the REST route `GET /agents/:id/operator-results`.
- */
-export interface OperatorResultSummary {
-  state: string;
-  inputTuples: number;
-  outputTuples: number;
-  inputPortShapes?: PortShape[];
-  outputColumns?: number;
-  error?: string;
-  warnings?: string[];
-  consoleLogCount?: number;
-  totalRowCount?: number;
-  sampleRecords?: Record<string, unknown>[];
-  resultStatistics?: Record<string, string>;
-}