Backlog/v12 soar command execution

backend/modules/soar/handler/command_ws.go

@@ -2,7 +2,6 @@ package handler
 
 import (
 	"context"
-	"encoding/json"
 	"net/http"
 	"strings"
 
@@ -83,7 +82,7 @@ type wsCommandRequest struct {
 }
 
 type wsMessage struct {
-	Type    string `json:"type"` // "output" | "error" | "done"
+	Type    string `json:"type"` // "output" | "error" | "ready"
 	Data    string `json:"data,omitempty"`
 	Message string `json:"message,omitempty"`
 }
@@ -119,12 +118,6 @@ func (h *CommandWSHandler) CommandStream(c *gin.Context) {
 	ctx, cancel := context.WithCancel(c.Request.Context())
 	defer cancel()
 
-	var req wsCommandRequest
-	if err := wsjson.Read(ctx, conn, &req); err != nil {
-		_ = conn.Close(websocket.StatusNormalClosure, "failed to read command")
-		return
-	}
-
 	if h.agentClient == nil {
 		_ = wsjson.Write(ctx, conn, wsMessage{Type: "error", Message: "agent manager unavailable"})
 		_ = conn.Close(websocket.StatusInternalError, "no agent client")
@@ -134,33 +127,51 @@ func (h *CommandWSHandler) CommandStream(c *gin.Context) {
 	// agent-manager validates it on ProcessCommandStream (offline/unknown → stream error,
 	// surfaced over the WS). Liveness for the UI comes from datasources, not from here.
 
-	command := req.Command
-	if h.variableUC != nil {
-		if interpolated, vErr := h.variableUC.InterpolateCommand(command); vErr != nil {
-			_ = catcher.Error("CommandStream: variable interpolation", vErr, nil)
-		} else {
-			command = interpolated
+	// Persistent session: read a command, execute, emit output + ready, wait for next.
+	// The socket stays open until the client disconnects or ctx is cancelled.
+	for {
+		var req wsCommandRequest
+		if err := wsjson.Read(ctx, conn, &req); err != nil {
+			return
 		}
-	}
 
-	cmd := &agent.UtmCommand{
-		AgentId:    agentID,
-		Command:    command,
-		ExecutedBy: loginFromCtx(c),
-		OriginType: req.OriginType,
-		OriginId:   req.OriginID,
-		Reason:     req.Reason,
-		Shell:      req.Shell,
+		command := req.Command
+		if h.variableUC != nil {
+			if interpolated, vErr := h.variableUC.InterpolateCommand(command); vErr != nil {
+				_ = catcher.Error("CommandStream: variable interpolation", vErr, nil)
+			} else {
+				command = interpolated
+			}
+		}
+
+		cmd := &agent.UtmCommand{
+			AgentId:    agentID,
+			Command:    command,
+			ExecutedBy: loginFromCtx(c),
+			OriginType: req.OriginType,
+			OriginId:   req.OriginID,
+			Reason:     req.Reason,
+			Shell:      req.Shell,
+		}
+		if !h.runCommand(ctx, conn, cmd) {
+			return
+		}
 	}
-	resultCh, errCh := h.agentClient.ProcessCommandStream(ctx, cmd)
+}
+
+// runCommand executes one command over a fresh gRPC stream, forwards output to
+// the WS, then emits "ready" so the client can send the next command. Returns
+// false if the WS is unusable (write failed / ctx cancelled) — caller must exit.
+func (h *CommandWSHandler) runCommand(ctx context.Context, conn *websocket.Conn, cmd *agent.UtmCommand) bool {
+	cmdCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+	resultCh, errCh := h.agentClient.ProcessCommandStream(cmdCtx, cmd)
 
 	for {
 		select {
 		case result, ok := <-resultCh:
 			if !ok {
-				_ = wsjson.Write(ctx, conn, wsMessage{Type: "done"})
-				_ = conn.Close(websocket.StatusNormalClosure, "stream complete")
-				return
+				return writeReady(ctx, conn)
 			}
 			output := result.GetResult()
 			if h.variableUC != nil {
@@ -170,26 +181,32 @@ func (h *CommandWSHandler) CommandStream(c *gin.Context) {
 					output = masked
 				}
 			}
-			msg := wsMessage{Type: "output", Data: output}
-			if writeErr := wsjson.Write(ctx, conn, msg); writeErr != nil {
-				cancel()
-				return
+			if err := wsjson.Write(ctx, conn, wsMessage{Type: "output", Data: output}); err != nil {
+				return false
 			}
+			// ponytail: protocol is 1-command→1-result (agent-manager sends one Send per
+			// command). First result IS end-of-stream — emit ready and return to caller,
+			// deferred cancel unblocks the gRPC recv. Upgrade if streaming ever becomes real.
+			return writeReady(ctx, conn)
 
 		case grpcErr, ok := <-errCh:
 			if !ok {
-				return
+				return writeReady(ctx, conn)
 			}
 			if grpcErr != nil {
 				_ = catcher.Error("CommandStream: grpc error", grpcErr, nil)
-				errMsg, _ := json.Marshal(wsMessage{Type: "error", Message: grpcErr.Error()})
-				_ = conn.Write(ctx, websocket.MessageText, errMsg)
-				_ = conn.Close(websocket.StatusInternalError, "grpc error")
+				if err := wsjson.Write(ctx, conn, wsMessage{Type: "error", Message: grpcErr.Error()}); err != nil {
+					return false
+				}
 			}
-			return
+			return writeReady(ctx, conn)
 
 		case <-ctx.Done():
-			return
+			return false
 		}
 	}
 }
+
+func writeReady(ctx context.Context, conn *websocket.Conn) bool {
+	return wsjson.Write(ctx, conn, wsMessage{Type: "ready"}) == nil
+}

frontend/src/features/datasources/components/AgentConsole.tsx

@@ -2,7 +2,7 @@ import { useEffect, useRef, useState } from 'react'
 import { useTranslation } from 'react-i18next'
 import { KeyRound, Loader2, TerminalSquare, X } from 'lucide-react'
 import { cn } from '@/shared/lib/utils'
-import { runAgentCommand } from '../services/agent-console'
+import { openAgentConsole, type ConsoleSession } from '../services/agent-console'
 import { VariablesManager } from './VariablesManager'
 import { ShellCommandInput } from './ShellCommandInput'
 
@@ -50,7 +50,7 @@ export function AgentConsole({
   const [histIdx, setHistIdx] = useState(-1)
   const [varsOpen, setVarsOpen] = useState(false)
   const [height, setHeight] = useState(() => clamp(Math.round(window.innerHeight * 0.42), 200, window.innerHeight - 80))
-  const cancelRef = useRef<(() => void) | null>(null)
+  const sessionRef = useRef<ConsoleSession | null>(null)
   const scrollRef = useRef<HTMLDivElement>(null)
   const inputRef = useRef<HTMLInputElement>(null)
 
@@ -72,9 +72,24 @@ export function AgentConsole({
 
   useEffect(() => {
     inputRef.current?.focus()
-    return () => cancelRef.current?.()
   }, [])
 
+  // Persistent console session: one WS per agent, reused for every command.
+  useEffect(() => {
+    const session = openAgentConsole(agentId, {
+      onOutput: (d) => append('out', d),
+      onError: (m) => append('err', m),
+      onReady: () => setRunning(false),
+      onClose: () => setRunning(false),
+    })
+    sessionRef.current = session
+    return () => {
+      sessionRef.current = null
+      session.close()
+    }
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [agentId])
+
   // Re-focus the inline prompt when a command finishes (the input remounts).
   useEffect(() => {
     if (!running) inputRef.current?.focus()
@@ -104,24 +119,13 @@ export function AgentConsole({
   const run = () => {
     const command = input.trim()
     if (!command || running || offline) return
+    if (!sessionRef.current) return
     append('cmd', `${shell}> ${command}`)
     setHistory((h) => [...h, command])
     setHistIdx(-1)
     setInput('')
     setRunning(true)
-    cancelRef.current = runAgentCommand(
-      agentId,
-      { command, shell },
-      {
-        onOutput: (d) => append('out', d),
-        onError: (m) => append('err', m),
-        onDone: () => {
-          setRunning(false)
-          cancelRef.current = null
-          inputRef.current?.focus()
-        },
-      },
-    )
+    sessionRef.current.send({ command, shell })
   }
 
   const onKey = (e: React.KeyboardEvent<HTMLInputElement>) => {

frontend/src/features/datasources/services/agent-console.ts

@@ -5,7 +5,7 @@ const API_URL = import.meta.env.VITE_API_URL || '/api/v1'
 
 /** A message streamed back from the agent command WebSocket. */
 interface ConsoleMessage {
-  type: 'output' | 'error' | 'done'
+  type: 'output' | 'error' | 'ready'
   data?: string
   message?: string
 }
@@ -18,45 +18,40 @@ export interface CommandPayload {
   reason?: string
 }
 
-export interface ConsoleHandlers {
+export interface SessionHandlers {
   onOutput: (data: string) => void
   onError: (message: string) => void
-  onDone: () => void
+  /** Server finished the current command and is ready for the next. */
+  onReady: () => void
+  /** Socket closed (network drop, server restart, or explicit close). */
+  onClose?: () => void
+}
+
+export interface ConsoleSession {
+  /** Send a command over the open socket. No-op if the socket isn't open. */
+  send: (payload: CommandPayload) => void
+  /** Close the socket. Idempotent. */
+  close: () => void
 }
 
 /**
- * Opens a WebSocket to the backend agent command stream
- * (GET /soar/ws/command/:agentId), sends ONE command, and streams its output.
- * The protocol is one command per connection — the server emits output chunks,
- * then a "done" message, then closes. Returns a cancel fn that closes the socket.
+ * Opens a persistent WebSocket to the backend agent command stream
+ * (GET /soar/ws/command/:agentId). One socket, many commands: each `send()`
+ * writes a JSON command frame, the server emits output frames, then a "ready"
+ * frame to unblock the next send. The socket stays open until close() or
+ * network failure.
  *
  * The JWT travels as a query param because browsers can't set WS headers; the
  * backend handler accepts `?token=` and verifies it.
  */
-export function runAgentCommand(agentId: string, payload: CommandPayload, h: ConsoleHandlers): () => void {
+export function openAgentConsole(agentId: string, h: SessionHandlers): ConsoleSession {
   const token = getStoredTokens()?.access_token ?? ''
   const proto = window.location.protocol === 'https:' ? 'wss' : 'ws'
   const url = `${proto}://${window.location.host}${API_URL}/soar/ws/command/${encodeURIComponent(agentId)}?token=${encodeURIComponent(token)}`
 
   const ws = new WebSocket(url)
-  let finished = false
-  const finish = () => {
-    if (!finished) {
-      finished = true
-      h.onDone()
-    }
-  }
+  let closed = false
 
-  ws.onopen = () => {
-    ws.send(
-      JSON.stringify({
-        originType: 'datasource',
-        originId: agentId,
-        reason: 'Interactive console',
-        ...payload,
-      }),
-    )
-  }
   ws.onmessage = (e) => {
     let msg: ConsoleMessage
     try {
@@ -65,22 +60,35 @@ export function runAgentCommand(agentId: string, payload: CommandPayload, h: Con
       return
     }
     if (msg.type === 'output') h.onOutput(msg.data ?? '')
-    else if (msg.type === 'error') {
-      h.onError(msg.message ?? i18n.t('datasources.console.cmdError'))
-      finish()
-    } else if (msg.type === 'done') finish()
+    else if (msg.type === 'error') h.onError(msg.message ?? i18n.t('datasources.console.cmdError'))
+    else if (msg.type === 'ready') h.onReady()
   }
-  ws.onerror = () => {
-    h.onError(i18n.t('datasources.console.connError'))
-    finish()
+  ws.onerror = () => h.onError(i18n.t('datasources.console.connError'))
+  ws.onclose = () => {
+    closed = true
+    h.onClose?.()
   }
-  ws.onclose = () => finish()
 
-  return () => {
-    try {
-      ws.close()
-    } catch {
-      /* noop */
-    }
+  return {
+    send: (payload) => {
+      if (ws.readyState !== WebSocket.OPEN) return
+      ws.send(
+        JSON.stringify({
+          originType: 'datasource',
+          originId: agentId,
+          reason: 'Interactive console',
+          ...payload,
+        }),
+      )
+    },
+    close: () => {
+      if (closed) return
+      closed = true
+      try {
+        ws.close()
+      } catch {
+        /* noop */
+      }
+    },
   }
 }