diff --git a/backend/modules/soar/handler/command_ws.go b/backend/modules/soar/handler/command_ws.go index 2bbe8cee4..e24aaa03b 100644 --- a/backend/modules/soar/handler/command_ws.go +++ b/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 +} diff --git a/frontend/src/features/datasources/components/AgentConsole.tsx b/frontend/src/features/datasources/components/AgentConsole.tsx index 21434e0a6..475a51138 100644 --- a/frontend/src/features/datasources/components/AgentConsole.tsx +++ b/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(null) const scrollRef = useRef(null) const inputRef = useRef(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) => { diff --git a/frontend/src/features/datasources/services/agent-console.ts b/frontend/src/features/datasources/services/agent-console.ts index fad79d304..4890c12b0 100644 --- a/frontend/src/features/datasources/services/agent-console.ts +++ b/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 */ + } + }, } }