resolve comments

Author: bobbai00

agent-service/src/server.ts

@@ -40,7 +40,7 @@ import type {
   ReActStep,
 } from "./types/agent";
 import { AgentState, OperatorResultSerializationMode } from "./types/agent";
-import type { WsClientRequest, WsServerMessage, WsServerSnapshotMessage, OperatorResultSummaryWs } from "./types/ws";
+import type { WsClientCommand, WsServerEvent, OperatorResultSummaryWs } from "./types/ws";
 
 const agentStore = new Map<string, TexeraAgent>();
 let agentCounter = 0;
@@ -325,7 +325,6 @@ const agentsRouter = new Elysia({ prefix: "/agents" })
       headId: stepId,
       steps: allSteps,
       workflowContent,
-      operatorResults: getOperatorResultSummaries(agent),
     });
 
     return {
@@ -434,7 +433,7 @@ function getOperatorResultSummaries(agent: TexeraAgent): Record<string, Operator
   return results;
 }
 
-function broadcastToAgent(agentId: string, message: WsServerMessage): void {
+function broadcastToAgent(agentId: string, message: WsServerEvent): void {
   const agent = agentStore.get(agentId);
   if (!agent) return;
 
@@ -474,7 +473,7 @@ export function buildApp() {
 
         agent.addWebsocket(ws);
 
-        const snapshotMessage: WsServerSnapshotMessage = {
+        const snapshotMessage: WsServerEvent = {
           type: "snapshot",
           state: agent.getState(),
           steps: agent.getAllSteps(),
@@ -492,19 +491,17 @@ export function buildApp() {
           return;
         }
 
-        let msg: WsClientRequest;
+        let msg: WsClientCommand;
         try {
-          msg = typeof messageData === "string" ? JSON.parse(messageData) : (messageData as WsClientRequest);
+          msg = typeof messageData === "string" ? JSON.parse(messageData) : (messageData as WsClientCommand);
         } catch {
           ws.send(JSON.stringify({ type: "error", error: "Invalid message format" }));
           return;
         }
 
-        if (msg.type === "command") {
-          if (msg.commandType === "stop") {
-            agent.stop();
-            broadcastToAgent(agentId, { type: "status", state: AgentState.STOPPING });
-          }
+        if (msg.type === "stop") {
+          agent.stop();
+          broadcastToAgent(agentId, { type: "status", state: AgentState.STOPPING });
           return;
         }
 
@@ -544,7 +541,13 @@ export function buildApp() {
             // from GENERATING after errors).
             broadcastToAgent(agentId, { type: "status", state: agent.getState() });
           }
+          return;
         }
+
+        // Frames are parsed from untrusted JSON; reject unknown discriminators
+        // explicitly instead of silently no-op'ing, so client/server mismatches
+        // are easy to diagnose.
+        ws.send(JSON.stringify({ type: "error", error: `Unknown message type: ${(msg as { type?: unknown }).type}` }));
       },
 
       close(ws) {
@@ -603,7 +606,7 @@ function printStartupMessage(app: ReturnType<typeof buildApp>) {
       console.log(`  WS     ${route.path}`);
     }
     console.log("         Send: { type: 'prompt', content: '...' }");
-    console.log("         Send: { type: 'command', commandType: 'stop' }");
+    console.log("         Send: { type: 'stop' }");
     console.log("         Recv: { type: 'snapshot' | 'step' | 'status' | 'error' | 'headChange', ... }");
   }
 

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

@@ -168,7 +168,7 @@ describe(`WS ${API}/agents/:id/react`, () => {
     await waitOpen(ws);
     await messages.waitFor(m => m.type === "snapshot");
 
-    ws.send(JSON.stringify({ type: "command", commandType: "stop" }));
+    ws.send(JSON.stringify({ type: "stop" }));
 
     const status = await messages.waitFor(m => m.type === "status");
     expect(status.state).toBe("STOPPING");
@@ -198,6 +198,18 @@ describe(`WS ${API}/agents/:id/react`, () => {
     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 === "snapshot");
+
+    ws.send(JSON.stringify({ type: "bogus" }));
+
+    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 () => {
     const id = await createAgent();
 

agent-service/src/types/index.ts

@@ -17,6 +17,7 @@
  * under the License.
  */
 
+export * from "./util";
 export * from "./workflow";
 export * from "./execution";
 export * from "./agent";

agent-service/src/types/util.ts

@@ -0,0 +1,27 @@
+/**
+ * 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.
+ */
+
+// Generic type-level utilities shared across the type modules.
+
+/**
+ * Builds a discriminated union from a `{ type -> payload }` map: each entry
+ * becomes `{ type: <key> } & <payload>`, then they are unioned. The `type` tag
+ * is derived from the map key rather than hand-written into each payload
+ */
+export type CustomUnionType<T> = { [K in keyof T]: { readonly type: K } & T[K] }[keyof T];

agent-service/src/types/ws/client.ts

@@ -18,32 +18,25 @@
  */
 
 // Client -> server WebSocket frames for this service's protocol
-// (`/agents/:id/react`). Modeled as a discriminated union on `type`, so each
-// request kind carries only its own fields.
 
-/** Shared discriminator base; every client request sets a unique `type`. */
-interface WsClientRequestBase {
-  type: "prompt" | "command";
-}
+import type { CustomUnionType } from "../util";
 
 /**
- * A user prompt for the agent to run. `messageSource` notes where it
- * originated (interactive chat vs. an operator feedback action).
+ * Send a prompt to agent to start the ReAct loop
  */
-export interface WsClientRequestPrompt extends WsClientRequestBase {
-  type: "prompt";
+export interface WsClientPromptCommand extends Readonly<{
   content: string;
   messageSource?: "chat" | "feedback";
-}
+}> {}
 
 /**
- * A control command. Today the only command stops the in-flight run; the
- * `commandType` discriminator leaves room for more commands later.
+ * Stop the agent's current ReAct loop
  */
-export interface WsClientRequestStopCommand extends WsClientRequestBase {
-  type: "command";
-  commandType: "stop";
-}
+export interface WsClientStopCommand extends Readonly<{}> {}
 
-/** Discriminated union of every client -> server frame. */
-export type WsClientRequest = WsClientRequestPrompt | WsClientRequestStopCommand;
+export type WsClientCommandTypeMap = {
+  prompt: WsClientPromptCommand;
+  stop: WsClientStopCommand;
+};
+
+export type WsClientCommand = CustomUnionType<WsClientCommandTypeMap>;

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

@@ -18,11 +18,10 @@
  */
 
 // Server -> client WebSocket frames for this service's protocol
-// (`/agents/:id/react`). Modeled as a discriminated union on `type`, so each
-// message kind declares exactly the fields it sends.
 
 import type { AgentState, ReActStep } from "../agent";
 import type { WorkflowContent } from "../workflow";
+import type { CustomUnionType } from "../util";
 
 /**
  * Wire projection of one operator's execution result, summarized for the
@@ -42,69 +41,54 @@ export interface OperatorResultSummaryWs {
   resultStatistics?: Record<string, string>;
 }
 
-/** Per-operator result summaries, keyed by operator id. */
-type OperatorResults = Record<string, OperatorResultSummaryWs>;
-
-/** Shared discriminator base; every server frame sets a unique `type`. */
-interface WsServerMessageBase {
-  type: "snapshot" | "step" | "status" | "error" | "headChange";
-}
-
 /**
  * Full state pushed once when a client connects: the agent's current lifecycle
  * state, the complete step list, and the HEAD pointer. Operator results are not
  * included — they are pulled on demand via `GET /operator-results`.
  */
-export interface WsServerSnapshotMessage extends WsServerMessageBase {
-  type: "snapshot";
+export interface WsServerSnapshotEvent extends Readonly<{
   state: AgentState;
   steps: ReActStep[];
   headId: string;
-}
+}> {}
 
-/**
- * A single ReAct step, streamed live as the agent runs.
- */
-export interface WsServerStepMessage extends WsServerMessageBase {
-  type: "step";
+/** A single ReAct step, streamed live as the agent runs. */
+export interface WsServerStepEvent extends Readonly<{
   step: ReActStep;
-}
+}> {}
 
 /**
  * An agent lifecycle transition (e.g. GENERATING when a run starts, the resting
  * state when it ends, STOPPING on stop).
  */
-export interface WsServerStatusMessage extends WsServerMessageBase {
-  type: "status";
+export interface WsServerStatusEvent extends Readonly<{
   state: AgentState;
-}
+}> {}
 
 /** An error surfaced to the client (agent not found, bad request, failed run). */
-export interface WsServerErrorMessage extends WsServerMessageBase {
-  type: "error";
+export interface WsServerErrorEvent extends Readonly<{
   error: string;
-}
+}> {}
 
 /**
  * Emitted after a checkout: HEAD moved, carrying the full step list and the
  * workflow snapshot at the new head.
  *
- * @deprecated Redundant and unused — the checkout flow that produces this frame
- * is unreachable in the product (nothing invokes the client's `checkoutStep()`).
- * Scheduled for removal (see #5930); do not build new code on it.
+ * @deprecated Redundant and unused. TODO: remove this message and related caller logics.
  */
-export interface WsServerHeadChangeMessage extends WsServerMessageBase {
-  type: "headChange";
+export interface WsServerHeadChangeEvent extends Readonly<{
   headId: string;
   steps: ReActStep[];
   workflowContent?: WorkflowContent;
-  operatorResults: OperatorResults;
-}
+}> {}
+
+export type WsServerEventTypeMap = {
+  snapshot: WsServerSnapshotEvent;
+  step: WsServerStepEvent;
+  status: WsServerStatusEvent;
+  error: WsServerErrorEvent;
+  headChange: WsServerHeadChangeEvent;
+};
 
 /** Discriminated union of every server -> client frame. */
-export type WsServerMessage =
-  | WsServerSnapshotMessage
-  | WsServerStepMessage
-  | WsServerStatusMessage
-  | WsServerErrorMessage
-  | WsServerHeadChangeMessage;
+export type WsServerEvent = CustomUnionType<WsServerEventTypeMap>;

frontend/src/app/workspace/service/agent/agent.service.ts

@@ -537,10 +537,6 @@ export class AgentService {
           };
           tracking.workflowSubject.next(workflow as Workflow);
         }
-        // Update operator results on HEAD change
-        if (message.operatorResults) {
-          this.updateOperatorResultSummaries(message.operatorResults);
-        }
         break;
 
       case "error":
@@ -952,7 +948,7 @@ export class AgentService {
     if (tracking?.websocket && tracking.websocket.readyState === WebSocket.OPEN) {
       // Send stop via WebSocket for immediate effect
       try {
-        tracking.websocket.send(JSON.stringify({ type: "command", commandType: "stop" }));
+        tracking.websocket.send(JSON.stringify({ type: "stop" }));
       } catch (error) {
         console.error("Failed to send stop command:", error);
       }