refactor(backend): split terminal_handler.go into pumps/sessions/messages (#1056)
* refactor(backend): split terminal_handler.go into pumps/sessions/messages terminal_handler.go was 1,380 lines (1.7x the 800-line cap). Split it into four sibling files in the same package, by responsibility: - terminal_handler.go (525 lines): HTTP upgrade, route dispatch, top-level handler entries, shell-label resolution. - terminal_pumps.go (725 lines): wsWriter, bridgeBinaryWebSockets, runTerminalBridge / runUserShellBridge and all input/resize/PTY-setup helpers. Extracts replayBufferedOutput from setupPtyAccessCommon to keep nestif under threshold. - terminal_sessions.go (133 lines): execution-readiness helpers (ensurePassthroughExecutionReady, waitForRemoteExecutionReadyWithTimeout, remoteExecutionWithClientReady, shouldUseWorkspaceShell) and related timeout constants. - terminal_messages.go (53 lines): protocol primitives (ResizePayload, isResizeCommand, resizeCommandByte, stripTerminalResponses), with a top-of-file doc comment describing the binary wire protocol. Backed off the "typed TerminalMessage envelope" half of the task: the current protocol is binary, not JSON-enveloped (terminal I/O is raw PTY bytes; the only JSON is ResizePayload, sent as a binary frame prefixed by 0x01). Wrapping PTY bytes in a JSON envelope would require base64 and break xterm.js AttachAddon for every client — the task brief explicitly authorizes shipping just the file split when the format would change. Adds terminal_messages_test.go pinning the wire format: exact JSON output of ResizePayload, the 0x01-prefix detection, and a full resize-frame round-trip. These fail loudly on any accidental change to the struct tags or prefix byte. Pure relocation otherwise: no signature changes, no behavior changes, no frontend impact. * refactor(backend): simplify detectInputSubmission with bytes.ContainsAny Addresses Claude bot's PR review on #1056: replace the manual loop + flag pattern with bytes.ContainsAny(data, "\n\r"). Behavior identical (detect CR or LF anywhere in the input). Also drops the now-redundant "Check for Enter key (CR or LF)" comment. * ci: retrigger CI to re-test known-flaky session-resume spec --------- Co-authored-by: Kandev Agent <agent@kandev.dev>
This commit is contained in:
João Salavisa
@@ -2,18 +2,12 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
gorillaws "github.com/gorilla/websocket"
|
||||
@@ -99,62 +93,6 @@ func checkWebSocketOrigin(r *http.Request) bool {
|
||||
return originHost == requestHost
|
||||
}
|
||||
|
||||
// resizeCommandByte is the binary protocol marker for resize messages.
|
||||
// First byte 0x01 indicates resize, followed by JSON {cols, rows}.
|
||||
const resizeCommandByte = 0x01
|
||||
|
||||
// isResizeCommand checks whether a binary frame is a resize command.
|
||||
// Resize messages use 0x01 as a prefix followed by a JSON object (starts with '{').
|
||||
// A bare 0x01 or 0x01 followed by non-JSON data is regular PTY input (e.g. Ctrl+A).
|
||||
func isResizeCommand(data []byte) bool {
|
||||
return len(data) >= 2 && data[0] == resizeCommandByte && data[1] == '{'
|
||||
}
|
||||
|
||||
const (
|
||||
passthroughReadyTimeout = 30 * time.Second
|
||||
remoteReadyProbeTimeout = 500 * time.Millisecond
|
||||
pollInterval = 500 * time.Millisecond
|
||||
)
|
||||
|
||||
// ResizePayload is the JSON payload for resize commands.
|
||||
type ResizePayload struct {
|
||||
Cols uint16 `json:"cols"`
|
||||
Rows uint16 `json:"rows"`
|
||||
}
|
||||
|
||||
// wsWriter wraps a gorilla WebSocket to implement process.DirectOutputWriter.
|
||||
// It writes binary frames to the WebSocket.
|
||||
type wsWriter struct {
|
||||
conn *gorillaws.Conn
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newWsWriter(conn *gorillaws.Conn) *wsWriter {
|
||||
return &wsWriter{conn: conn}
|
||||
}
|
||||
|
||||
func (w *wsWriter) Write(p []byte) (n int, err error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
if w.closed {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
|
||||
if err := w.conn.WriteMessage(gorillaws.BinaryMessage, p); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (w *wsWriter) Close() error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
type terminalRoute struct {
|
||||
kind string
|
||||
id string
|
||||
@@ -240,83 +178,6 @@ func (h *TerminalHandler) handleAgentTerminalRoute(c *gin.Context, sessionID str
|
||||
h.handleAgentPassthroughWS(c, sessionID, interactiveRunner)
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) shouldUseWorkspaceShell(ctx context.Context, sessionID string) bool {
|
||||
// Delegates to lifecycle manager which checks both in-memory execution
|
||||
// and database records for containerized/remote executors.
|
||||
return h.lifecycleMgr.ShouldUseContainerShell(ctx, sessionID)
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) waitForRemoteExecutionReadyWithTimeout(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
timeout time.Duration,
|
||||
) (*lifecycle.AgentExecution, bool) {
|
||||
deadline := time.Now().Add(timeout)
|
||||
|
||||
h.logger.Info("waiting for remote execution to become ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", timeout))
|
||||
|
||||
for {
|
||||
if execution, ok := h.remoteExecutionWithClientReady(ctx, sessionID, deadline); ok {
|
||||
return execution, true
|
||||
}
|
||||
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
h.logger.Warn("timed out waiting for remote execution readiness",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", timeout))
|
||||
return nil, false
|
||||
}
|
||||
|
||||
wait := pollInterval
|
||||
if remaining < wait {
|
||||
wait = remaining
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return nil, false
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) remoteExecutionWithClientReady(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
deadline time.Time,
|
||||
) (*lifecycle.AgentExecution, bool) {
|
||||
execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID)
|
||||
if !exists || execution.GetAgentCtlClient() == nil {
|
||||
return nil, false
|
||||
}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
probeTimeout := remoteReadyProbeTimeout
|
||||
if remaining < probeTimeout {
|
||||
probeTimeout = remaining
|
||||
}
|
||||
probeCtx, cancel := context.WithTimeout(ctx, probeTimeout)
|
||||
defer cancel()
|
||||
if err := execution.GetAgentCtlClient().Health(probeCtx); err != nil {
|
||||
h.logger.Debug("remote execution agentctl client not ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID),
|
||||
zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
h.logger.Info("remote execution became ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID))
|
||||
return execution, true
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) handleEnvironmentUserShellWS(c *gin.Context, environmentID, terminalID string) {
|
||||
execution, err := h.lifecycleMgr.GetOrEnsureExecutionForEnvironment(c.Request.Context(), environmentID)
|
||||
if err != nil {
|
||||
@@ -434,58 +295,6 @@ func (h *TerminalHandler) handleRemoteUserShellWS(
|
||||
h.bridgeBinaryWebSockets(clientConn, agentctlConn, terminalID)
|
||||
}
|
||||
|
||||
// bridgeBinaryWebSockets forwards binary WebSocket frames bidirectionally between
|
||||
// a client connection and an agentctl connection until either side disconnects.
|
||||
func (h *TerminalHandler) bridgeBinaryWebSockets(clientConn, agentctlConn *gorillaws.Conn, terminalID string) {
|
||||
done := make(chan struct{})
|
||||
|
||||
// agentctl → client (shell output)
|
||||
go func() {
|
||||
defer close(done)
|
||||
for {
|
||||
msgType, data, err := agentctlConn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("remote terminal shell upstream read error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := clientConn.WriteMessage(msgType, data); err != nil {
|
||||
h.logger.Debug("remote terminal shell client write error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// client → agentctl (input + resize)
|
||||
go func() {
|
||||
for {
|
||||
msgType, data, err := clientConn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("remote terminal shell client read error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
_ = agentctlConn.Close()
|
||||
return
|
||||
}
|
||||
if err := agentctlConn.WriteMessage(msgType, data); err != nil {
|
||||
h.logger.Debug("remote terminal shell upstream write error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
<-done
|
||||
}
|
||||
|
||||
// handleAgentPassthroughWS handles WebSocket connections for agent passthrough terminals.
|
||||
// It ensures a passthrough execution exists, upgrades to WebSocket, and bridges I/O.
|
||||
func (h *TerminalHandler) handleAgentPassthroughWS(
|
||||
@@ -550,364 +359,6 @@ func (h *TerminalHandler) handleAgentPassthroughWS(
|
||||
h.runTerminalBridge(conn, sessionID, processID, interactiveRunner, wsw)
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) ensurePassthroughExecutionReady(ctx context.Context, sessionID string) (*lifecycle.AgentExecution, error) {
|
||||
execution, err := h.lifecycleMgr.EnsurePassthroughExecution(ctx, sessionID)
|
||||
if err == nil {
|
||||
return execution, nil
|
||||
}
|
||||
if !errors.Is(err, lifecycle.ErrSessionWorkspaceNotReady) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
h.logger.Info("waiting for passthrough workspace to become ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", passthroughReadyTimeout))
|
||||
|
||||
timeoutCtx, cancel := context.WithTimeout(ctx, passthroughReadyTimeout)
|
||||
defer cancel()
|
||||
|
||||
ticker := time.NewTicker(pollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timeoutCtx.Done():
|
||||
return nil, fmt.Errorf("%w: timed out after %s", lifecycle.ErrSessionWorkspaceNotReady, passthroughReadyTimeout)
|
||||
case <-ticker.C:
|
||||
execution, err := h.lifecycleMgr.EnsurePassthroughExecution(timeoutCtx, sessionID)
|
||||
if err == nil {
|
||||
h.logger.Info("passthrough workspace became ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID))
|
||||
return execution, nil
|
||||
}
|
||||
if !errors.Is(err, lifecycle.ErrSessionWorkspaceNotReady) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// runTerminalBridge handles WebSocket I/O and manages the PTY connection.
|
||||
// It waits for the first resize to trigger process start, then sets up bidirectional I/O.
|
||||
func (h *TerminalHandler) runTerminalBridge(
|
||||
conn *gorillaws.Conn,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
) {
|
||||
var ptyWriter io.Writer
|
||||
var directOutputSet bool
|
||||
|
||||
defer func() {
|
||||
// Clean up: clear direct output at session level (process may have been restarted with new ID)
|
||||
interactiveRunner.ClearDirectOutputBySession(sessionID)
|
||||
_ = wsw.Close()
|
||||
_ = conn.Close()
|
||||
|
||||
h.logger.Info("terminal WebSocket disconnected",
|
||||
zap.String("session_id", sessionID))
|
||||
}()
|
||||
|
||||
// Read from WebSocket and handle input/resize
|
||||
for {
|
||||
messageType, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("WebSocket read error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if messageType != gorillaws.BinaryMessage && messageType != gorillaws.TextMessage {
|
||||
continue
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if isResizeCommand(data) {
|
||||
newID, newDirect := h.handleResizeCommand(data[1:], sessionID, processID, interactiveRunner, wsw, directOutputSet, &ptyWriter)
|
||||
processID = newID
|
||||
directOutputSet = newDirect
|
||||
continue
|
||||
}
|
||||
|
||||
newID, newDirect := h.handleTerminalInput(data, sessionID, processID, interactiveRunner, wsw, &ptyWriter, directOutputSet)
|
||||
processID = newID
|
||||
directOutputSet = newDirect
|
||||
}
|
||||
}
|
||||
|
||||
// handleResizeCommand processes a resize message in the terminal bridge loop.
|
||||
// Returns the (possibly updated) processID and directOutputSet flag.
|
||||
func (h *TerminalHandler) handleResizeCommand(
|
||||
payload []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
directOutputSet bool,
|
||||
ptyWriter *io.Writer,
|
||||
) (string, bool) {
|
||||
// Set up PTY access BEFORE handling resize if not already done.
|
||||
// This ensures buffered output is sent to the terminal before the resize
|
||||
// triggers a TUI redraw (which would overwrite the scrollback).
|
||||
if !directOutputSet {
|
||||
if h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
|
||||
// Handle resize — may return updated processID if process was restarted.
|
||||
// Falls back to session-level resize when the old process is gone.
|
||||
newProcessID := h.handleResize(payload, sessionID, processID, interactiveRunner)
|
||||
if newProcessID != processID {
|
||||
processID = newProcessID
|
||||
directOutputSet = false
|
||||
*ptyWriter = nil
|
||||
}
|
||||
|
||||
// Retry setupPtyAccess after resize — process may have been lazily started
|
||||
// by the resize (deferred-start), so the PTY is now available.
|
||||
if !directOutputSet {
|
||||
if h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// handleTerminalInput processes regular (non-resize) input in the terminal bridge loop.
|
||||
// Returns the (possibly updated) processID and directOutputSet flag.
|
||||
func (h *TerminalHandler) handleTerminalInput(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
if *ptyWriter != nil {
|
||||
return h.writeToReadyPty(data, sessionID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
return h.writeToUnreadyPty(data, sessionID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
|
||||
// writeToReadyPty writes input to an already-established PTY writer.
|
||||
// If writing fails it resets state and discovers the new process.
|
||||
func (h *TerminalHandler) writeToReadyPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(err))
|
||||
|
||||
// Reset state and discover the new process immediately.
|
||||
// Without this, the bridge would be stuck with the old processID
|
||||
// until a resize message arrives (which may never come after restart).
|
||||
*ptyWriter = nil
|
||||
directOutputSet = false
|
||||
processID = h.discoverProcessBySession(sessionID, processID, interactiveRunner)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
// Detect Enter key (user submitted input) - notify lifecycle manager
|
||||
h.detectInputSubmission(sessionID, data)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// writeToUnreadyPty attempts to set up PTY access and then write input.
|
||||
// If setup fails it tries to discover the current process by session.
|
||||
func (h *TerminalHandler) writeToUnreadyPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
// PTY not ready — try setup. If it fails, processID may be stale
|
||||
// (process was restarted with a new ID). Discover the current process
|
||||
// by session before giving up.
|
||||
if !h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
newID := h.discoverProcessBySession(sessionID, processID, interactiveRunner)
|
||||
if newID != processID {
|
||||
processID = newID
|
||||
h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter)
|
||||
}
|
||||
}
|
||||
if *ptyWriter == nil {
|
||||
h.logger.Debug("PTY not ready, dropping input",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Int("bytes", len(data)))
|
||||
return processID, directOutputSet
|
||||
}
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after setup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
return processID, directOutputSet
|
||||
}
|
||||
h.detectInputSubmission(sessionID, data)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// detectInputSubmission checks if the input contains Enter key and notifies
|
||||
// the lifecycle manager to update the execution state to Running.
|
||||
func (h *TerminalHandler) detectInputSubmission(sessionID string, data []byte) {
|
||||
// Check for Enter key (CR or LF)
|
||||
hasEnter := false
|
||||
for _, b := range data {
|
||||
if b == '\n' || b == '\r' {
|
||||
hasEnter = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if hasEnter {
|
||||
// Notify lifecycle manager that user submitted input
|
||||
if err := h.lifecycleMgr.MarkPassthroughRunning(sessionID); err != nil {
|
||||
// Log at debug level - may fail if session ended or not in passthrough mode
|
||||
h.logger.Debug("failed to mark passthrough as running",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// setupPtyAccess attempts to get PTY writer and set up direct output.
|
||||
// Returns true if successful.
|
||||
func (h *TerminalHandler) setupPtyAccess(
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
getWriter := func() (io.Writer, error) {
|
||||
return interactiveRunner.GetPtyWriter(processID)
|
||||
}
|
||||
return h.setupPtyAccessCommon(sessionID, processID, interactiveRunner, wsw, ptyWriter, getWriter)
|
||||
}
|
||||
|
||||
// handleResize processes a resize command from the WebSocket.
|
||||
// Uses process-specific resize to avoid targeting the wrong process when multiple
|
||||
// processes exist for the same session (e.g., passthrough + user shell).
|
||||
// If the process was restarted (old process gone), falls back to session-level resize
|
||||
// and returns the new process ID.
|
||||
func (h *TerminalHandler) handleResize(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) (effectiveProcessID string) {
|
||||
effectiveProcessID = processID
|
||||
|
||||
var resize ResizePayload
|
||||
if err := json.Unmarshal(data, &resize); err != nil {
|
||||
h.logger.Warn("failed to parse resize command",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
if resize.Cols == 0 || resize.Rows == 0 {
|
||||
h.logger.Warn("invalid resize dimensions",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
// Resize the specific process by ID to avoid targeting the wrong process.
|
||||
// If the process was restarted (old ID gone), fall back to session-level resize.
|
||||
resizeErr := interactiveRunner.ResizeByProcessID(processID, resize.Cols, resize.Rows)
|
||||
if resizeErr == nil {
|
||||
h.logger.Debug("PTY resized",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
}
|
||||
|
||||
if resizeErr != nil {
|
||||
h.logger.Debug("process-specific resize failed, trying session-level fallback",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(resizeErr))
|
||||
effectiveProcessID = h.resizeBySessionFallback(sessionID, processID, resize.Cols, resize.Rows, interactiveRunner)
|
||||
}
|
||||
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
// discoverProcessBySession looks up the current passthrough process for a session.
|
||||
// Returns the new processID if found and different from oldProcessID.
|
||||
// This is used when the bridge has a stale processID after a process restart.
|
||||
func (h *TerminalHandler) discoverProcessBySession(
|
||||
sessionID string,
|
||||
oldProcessID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) string {
|
||||
_, newProcID, _ := interactiveRunner.GetPtyWriterBySession(sessionID)
|
||||
if newProcID != "" && newProcID != oldProcessID {
|
||||
h.logger.Info("discovered new process by session lookup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("old_process_id", oldProcessID),
|
||||
zap.String("new_process_id", newProcID))
|
||||
return newProcID
|
||||
}
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// resizeBySessionFallback attempts a session-level resize when the process-specific
|
||||
// resize failed (e.g., process was restarted). Returns the new processID if found.
|
||||
func (h *TerminalHandler) resizeBySessionFallback(
|
||||
sessionID string,
|
||||
oldProcessID string,
|
||||
cols, rows uint16,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) string {
|
||||
sessionErr := interactiveRunner.ResizeBySession(sessionID, cols, rows)
|
||||
if sessionErr != nil {
|
||||
h.logger.Warn("failed to resize PTY (both process and session)",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", oldProcessID),
|
||||
zap.Uint16("cols", cols),
|
||||
zap.Uint16("rows", rows),
|
||||
zap.Error(sessionErr))
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// Session-level resize succeeded — look up the new process ID
|
||||
_, newProcID, _ := interactiveRunner.GetPtyWriterBySession(sessionID)
|
||||
if newProcID != "" && newProcID != oldProcessID {
|
||||
h.logger.Info("resize fallback succeeded, process ID updated",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("old_process_id", oldProcessID),
|
||||
zap.String("new_process_id", newProcID))
|
||||
return newProcID
|
||||
}
|
||||
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// shellScriptInfo holds the label and initial command resolved from a script.
|
||||
type shellScriptInfo struct {
|
||||
label string
|
||||
@@ -1072,310 +523,3 @@ func (h *TerminalHandler) handleUserShellWS(
|
||||
wsw := newWsWriter(conn)
|
||||
h.runUserShellBridge(conn, sessionID, scopeID, terminalID, processID, interactiveRunner, wsw)
|
||||
}
|
||||
|
||||
// runUserShellBridge handles WebSocket I/O for user shell terminals.
|
||||
func (h *TerminalHandler) runUserShellBridge(
|
||||
conn *gorillaws.Conn,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
) {
|
||||
var ptyWriter io.Writer
|
||||
var directOutputSet bool
|
||||
|
||||
defer func() {
|
||||
// Clean up WebSocket resources but DON'T stop the process
|
||||
// The process should only be stopped via explicit user_shell.stop message from frontend
|
||||
// This allows reconnection after React remounts or temporary disconnects
|
||||
interactiveRunner.ClearUserShellDirectOutput(scopeID, terminalID)
|
||||
|
||||
_ = wsw.Close()
|
||||
_ = conn.Close()
|
||||
|
||||
h.logger.Info("user shell WebSocket disconnected (process still running for potential reconnect)",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID))
|
||||
}()
|
||||
|
||||
// Read from WebSocket and handle input/resize
|
||||
for {
|
||||
messageType, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("WebSocket read error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if messageType != gorillaws.BinaryMessage && messageType != gorillaws.TextMessage {
|
||||
continue
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if isResizeCommand(data) {
|
||||
directOutputSet = h.handleUserShellResizeCommand(
|
||||
data[1:], sessionID, scopeID, terminalID, processID,
|
||||
interactiveRunner, wsw, directOutputSet, &ptyWriter,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
directOutputSet = h.handleUserShellInput(
|
||||
data, sessionID, scopeID, terminalID, processID,
|
||||
interactiveRunner, wsw, &ptyWriter, directOutputSet,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// handleUserShellResizeCommand processes a resize message in the user shell bridge loop.
|
||||
// Returns the updated directOutputSet flag.
|
||||
func (h *TerminalHandler) handleUserShellResizeCommand(
|
||||
payload []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
directOutputSet bool,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
// Set up PTY access BEFORE handling resize if not already done
|
||||
if !directOutputSet {
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
h.handleUserShellResize(payload, sessionID, scopeID, terminalID, interactiveRunner)
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// handleUserShellInput processes regular (non-resize) input in the user shell bridge loop.
|
||||
// Returns the updated directOutputSet flag.
|
||||
func (h *TerminalHandler) handleUserShellInput(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if *ptyWriter != nil {
|
||||
return h.writeToReadyUserShellPty(data, sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
return h.writeToUnreadyUserShellPty(data, sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
|
||||
// writeToReadyUserShellPty writes input to an established user shell PTY.
|
||||
// On failure it attempts to re-establish PTY access and retries the write.
|
||||
func (h *TerminalHandler) writeToReadyUserShellPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
// Try to re-establish PTY access
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after reconnect",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// writeToUnreadyUserShellPty sets up PTY access and writes input for a user shell.
|
||||
// Drops the input if setup fails.
|
||||
func (h *TerminalHandler) writeToUnreadyUserShellPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after setup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
} else {
|
||||
h.logger.Debug("PTY not ready, dropping input",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Int("bytes", len(data)))
|
||||
}
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// setupUserShellPtyAccess sets up PTY access for a user shell terminal.
|
||||
func (h *TerminalHandler) setupUserShellPtyAccess(
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
getWriter := func() (io.Writer, error) {
|
||||
writer, _, err := interactiveRunner.GetUserShellPtyWriter(scopeID, terminalID)
|
||||
return writer, err
|
||||
}
|
||||
return h.setupPtyAccessCommon(sessionID, processID, interactiveRunner, wsw, ptyWriter, getWriter)
|
||||
}
|
||||
|
||||
// terminalResponsePattern matches terminal query responses that should not be replayed.
|
||||
// These are responses from the terminal emulator to queries sent by the shell, which
|
||||
// appear as PTY output and get captured in the ring buffer. On replay they render as
|
||||
// visible garbage text because they are response sequences, not display sequences.
|
||||
//
|
||||
// Matched sequences:
|
||||
// - OSC 11 response: \x1b]11;rgb:XXXX/XXXX/XXXX ST (background color, ST = ESC\ or BEL)
|
||||
// - DA1 response: \x1b[?<params>c (device attributes)
|
||||
// - CPR response: \x1b[<row>;<col>R or \x1b[<row>R (cursor position report)
|
||||
var terminalResponsePattern = regexp.MustCompile(
|
||||
`\x1b]11;rgb:[0-9a-fA-F/]+(?:\x1b\\|\x07)` + // OSC 11 (ESC\ or BEL terminator)
|
||||
`|\x1b\[\?[0-9;]*c` + // DA1 response
|
||||
`|\x1b\[\d+(?:;\d+)?R`, // CPR response
|
||||
)
|
||||
|
||||
// stripTerminalResponses removes terminal query response sequences from data.
|
||||
// Used to clean buffered PTY output before replaying it on reconnect.
|
||||
func stripTerminalResponses(data []byte) []byte {
|
||||
return terminalResponsePattern.ReplaceAll(data, nil)
|
||||
}
|
||||
|
||||
// setupPtyAccessCommon is the shared implementation for setupPtyAccess and setupUserShellPtyAccess.
|
||||
// It gets a PTY writer via getWriter, sends buffered output, and sets up direct output.
|
||||
func (h *TerminalHandler) setupPtyAccessCommon(
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
getWriter func() (io.Writer, error),
|
||||
) bool {
|
||||
writer, err := getWriter()
|
||||
if err != nil {
|
||||
h.logger.Debug("PTY writer not ready yet",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID))
|
||||
return false
|
||||
}
|
||||
|
||||
// Send buffered output to the terminal before switching to direct mode.
|
||||
// This provides scrollback history when reconnecting to an existing session.
|
||||
if chunks, ok := interactiveRunner.GetBuffer(processID); ok && len(chunks) > 0 {
|
||||
var combined bytes.Buffer
|
||||
for _, chunk := range chunks {
|
||||
combined.WriteString(chunk.Data)
|
||||
}
|
||||
if combined.Len() > 0 {
|
||||
cleaned := stripTerminalResponses(combined.Bytes())
|
||||
h.logger.Debug("sending buffered output to terminal",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Int("chunks", len(chunks)),
|
||||
zap.Int("total_bytes", combined.Len()),
|
||||
zap.Int("cleaned_bytes", len(cleaned)))
|
||||
if len(cleaned) > 0 {
|
||||
if _, writeErr := wsw.Write(cleaned); writeErr != nil {
|
||||
h.logger.Warn("failed to send buffered output",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(writeErr))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set up direct output
|
||||
if err := interactiveRunner.SetDirectOutput(processID, wsw); err != nil {
|
||||
h.logger.Error("failed to set direct output",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
*ptyWriter = writer
|
||||
|
||||
h.logger.Info("PTY access established",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID))
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// handleUserShellResize processes resize commands for user shell terminals.
|
||||
func (h *TerminalHandler) handleUserShellResize(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) {
|
||||
var resize ResizePayload
|
||||
if err := json.Unmarshal(data, &resize); err != nil {
|
||||
h.logger.Warn("failed to parse resize command",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
if resize.Cols == 0 || resize.Rows == 0 {
|
||||
h.logger.Warn("invalid resize dimensions",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
return
|
||||
}
|
||||
|
||||
if err := interactiveRunner.ResizeUserShell(scopeID, terminalID, resize.Cols, resize.Rows); err != nil {
|
||||
h.logger.Warn("failed to resize user shell PTY",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows),
|
||||
zap.Error(err))
|
||||
} else {
|
||||
h.logger.Debug("user shell PTY resized",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
package websocket
|
||||
|
||||
import "regexp"
|
||||
|
||||
// Terminal WebSocket wire protocol
|
||||
// ---------------------------------
|
||||
// Frames are binary WebSocket messages. The first byte distinguishes the
|
||||
// frame kind:
|
||||
//
|
||||
// - 0x01 + '{' ... '}' → JSON resize command (ResizePayload)
|
||||
// - anything else → raw PTY input/output bytes
|
||||
//
|
||||
// Output to the client is always raw PTY bytes — never wrapped. Wrapping
|
||||
// PTY data in a JSON envelope would require base64 and break the xterm.js
|
||||
// AttachAddon, so the protocol stays binary on purpose.
|
||||
|
||||
// resizeCommandByte is the binary protocol marker for resize messages.
|
||||
// First byte 0x01 indicates resize, followed by JSON {cols, rows}.
|
||||
const resizeCommandByte = 0x01
|
||||
|
||||
// isResizeCommand checks whether a binary frame is a resize command.
|
||||
// Resize messages use 0x01 as a prefix followed by a JSON object (starts with '{').
|
||||
// A bare 0x01 or 0x01 followed by non-JSON data is regular PTY input (e.g. Ctrl+A).
|
||||
func isResizeCommand(data []byte) bool {
|
||||
return len(data) >= 2 && data[0] == resizeCommandByte && data[1] == '{'
|
||||
}
|
||||
|
||||
// ResizePayload is the JSON payload for resize commands.
|
||||
type ResizePayload struct {
|
||||
Cols uint16 `json:"cols"`
|
||||
Rows uint16 `json:"rows"`
|
||||
}
|
||||
|
||||
// terminalResponsePattern matches terminal query responses that should not be replayed.
|
||||
// These are responses from the terminal emulator to queries sent by the shell, which
|
||||
// appear as PTY output and get captured in the ring buffer. On replay they render as
|
||||
// visible garbage text because they are response sequences, not display sequences.
|
||||
//
|
||||
// Matched sequences:
|
||||
// - OSC 11 response: \x1b]11;rgb:XXXX/XXXX/XXXX ST (background color, ST = ESC\ or BEL)
|
||||
// - DA1 response: \x1b[?<params>c (device attributes)
|
||||
// - CPR response: \x1b[<row>;<col>R or \x1b[<row>R (cursor position report)
|
||||
var terminalResponsePattern = regexp.MustCompile(
|
||||
`\x1b]11;rgb:[0-9a-fA-F/]+(?:\x1b\\|\x07)` + // OSC 11 (ESC\ or BEL terminator)
|
||||
`|\x1b\[\?[0-9;]*c` + // DA1 response
|
||||
`|\x1b\[\d+(?:;\d+)?R`, // CPR response
|
||||
)
|
||||
|
||||
// stripTerminalResponses removes terminal query response sequences from data.
|
||||
// Used to clean buffered PTY output before replaying it on reconnect.
|
||||
func stripTerminalResponses(data []byte) []byte {
|
||||
return terminalResponsePattern.ReplaceAll(data, nil)
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestResizePayloadWireFormat locks in the exact JSON encoding the frontend
|
||||
// expects. Any change to ResizePayload field tags would break every connected
|
||||
// terminal client. If this fails after a struct edit, the wire format changed
|
||||
// and the frontend would need to be updated in lockstep.
|
||||
func TestResizePayloadWireFormat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
payload ResizePayload
|
||||
want string
|
||||
}{
|
||||
{name: "typical 80x24", payload: ResizePayload{Cols: 80, Rows: 24}, want: `{"cols":80,"rows":24}`},
|
||||
{name: "zero", payload: ResizePayload{Cols: 0, Rows: 0}, want: `{"cols":0,"rows":0}`},
|
||||
{name: "max uint16", payload: ResizePayload{Cols: 65535, Rows: 65535}, want: `{"cols":65535,"rows":65535}`},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got, err := json.Marshal(tt.payload)
|
||||
if err != nil {
|
||||
t.Fatalf("marshal: %v", err)
|
||||
}
|
||||
if string(got) != tt.want {
|
||||
t.Fatalf("encoded = %q, want %q", got, tt.want)
|
||||
}
|
||||
|
||||
// Round-trip: server must be able to decode what it (and the
|
||||
// frontend) emit. Field names are lowercase JSON tags.
|
||||
var decoded ResizePayload
|
||||
if err := json.Unmarshal(got, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal: %v", err)
|
||||
}
|
||||
if decoded != tt.payload {
|
||||
t.Fatalf("round-trip = %+v, want %+v", decoded, tt.payload)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestIsResizeCommandWireFormat documents the binary protocol prefix used to
|
||||
// distinguish resize commands from raw PTY input. Changing either resizeCommandByte
|
||||
// or the leading-'{' heuristic would silently route resize frames to the PTY (and
|
||||
// vice versa) for any connected client.
|
||||
func TestIsResizeCommandWireFormat(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
want bool
|
||||
}{
|
||||
{name: "resize prefix + json", data: []byte{0x01, '{', '}'}, want: true},
|
||||
{name: "resize prefix with realistic payload", data: append([]byte{0x01}, []byte(`{"cols":80,"rows":24}`)...), want: true},
|
||||
{name: "bare 0x01 (Ctrl+A)", data: []byte{0x01}, want: false},
|
||||
{name: "0x01 followed by non-brace", data: []byte{0x01, 'X'}, want: false},
|
||||
{name: "regular input", data: []byte("ls\n"), want: false},
|
||||
{name: "empty", data: []byte{}, want: false},
|
||||
{name: "json without prefix", data: []byte(`{"cols":80,"rows":24}`), want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
if got := isResizeCommand(tt.data); got != tt.want {
|
||||
t.Fatalf("isResizeCommand(%q) = %v, want %v", tt.data, got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestResizeFramePayloadParseable verifies the full wire frame (prefix byte +
|
||||
// JSON body) round-trips through the handler's parsing path: strip the prefix,
|
||||
// json.Unmarshal the remainder into ResizePayload.
|
||||
func TestResizeFramePayloadParseable(t *testing.T) {
|
||||
frame := []byte{0x01}
|
||||
frame = append(frame, []byte(`{"cols":120,"rows":40}`)...)
|
||||
|
||||
if !isResizeCommand(frame) {
|
||||
t.Fatalf("frame not detected as resize: %q", frame)
|
||||
}
|
||||
if !bytes.HasPrefix(frame, []byte{resizeCommandByte}) {
|
||||
t.Fatalf("frame missing resize prefix byte")
|
||||
}
|
||||
|
||||
var resize ResizePayload
|
||||
if err := json.Unmarshal(frame[1:], &resize); err != nil {
|
||||
t.Fatalf("unmarshal frame payload: %v", err)
|
||||
}
|
||||
if resize.Cols != 120 || resize.Rows != 40 {
|
||||
t.Fatalf("parsed = %+v, want cols=120 rows=40", resize)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,712 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
gorillaws "github.com/gorilla/websocket"
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/kandev/kandev/internal/agentctl/server/process"
|
||||
)
|
||||
|
||||
// wsWriter wraps a gorilla WebSocket to implement process.DirectOutputWriter.
|
||||
// It writes binary frames to the WebSocket.
|
||||
type wsWriter struct {
|
||||
conn *gorillaws.Conn
|
||||
mu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newWsWriter(conn *gorillaws.Conn) *wsWriter {
|
||||
return &wsWriter{conn: conn}
|
||||
}
|
||||
|
||||
func (w *wsWriter) Write(p []byte) (n int, err error) {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
|
||||
if w.closed {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
|
||||
if err := w.conn.WriteMessage(gorillaws.BinaryMessage, p); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
return len(p), nil
|
||||
}
|
||||
|
||||
func (w *wsWriter) Close() error {
|
||||
w.mu.Lock()
|
||||
defer w.mu.Unlock()
|
||||
w.closed = true
|
||||
return nil
|
||||
}
|
||||
|
||||
// bridgeBinaryWebSockets forwards binary WebSocket frames bidirectionally between
|
||||
// a client connection and an agentctl connection until either side disconnects.
|
||||
func (h *TerminalHandler) bridgeBinaryWebSockets(clientConn, agentctlConn *gorillaws.Conn, terminalID string) {
|
||||
done := make(chan struct{})
|
||||
|
||||
// agentctl → client (shell output)
|
||||
go func() {
|
||||
defer close(done)
|
||||
for {
|
||||
msgType, data, err := agentctlConn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("remote terminal shell upstream read error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
if err := clientConn.WriteMessage(msgType, data); err != nil {
|
||||
h.logger.Debug("remote terminal shell client write error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// client → agentctl (input + resize)
|
||||
go func() {
|
||||
for {
|
||||
msgType, data, err := clientConn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("remote terminal shell client read error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
_ = agentctlConn.Close()
|
||||
return
|
||||
}
|
||||
if err := agentctlConn.WriteMessage(msgType, data); err != nil {
|
||||
h.logger.Debug("remote terminal shell upstream write error",
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
<-done
|
||||
}
|
||||
|
||||
// runTerminalBridge handles WebSocket I/O and manages the PTY connection.
|
||||
// It waits for the first resize to trigger process start, then sets up bidirectional I/O.
|
||||
func (h *TerminalHandler) runTerminalBridge(
|
||||
conn *gorillaws.Conn,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
) {
|
||||
var ptyWriter io.Writer
|
||||
var directOutputSet bool
|
||||
|
||||
defer func() {
|
||||
// Clean up: clear direct output at session level (process may have been restarted with new ID)
|
||||
interactiveRunner.ClearDirectOutputBySession(sessionID)
|
||||
_ = wsw.Close()
|
||||
_ = conn.Close()
|
||||
|
||||
h.logger.Info("terminal WebSocket disconnected",
|
||||
zap.String("session_id", sessionID))
|
||||
}()
|
||||
|
||||
// Read from WebSocket and handle input/resize
|
||||
for {
|
||||
messageType, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("WebSocket read error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if messageType != gorillaws.BinaryMessage && messageType != gorillaws.TextMessage {
|
||||
continue
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if isResizeCommand(data) {
|
||||
newID, newDirect := h.handleResizeCommand(data[1:], sessionID, processID, interactiveRunner, wsw, directOutputSet, &ptyWriter)
|
||||
processID = newID
|
||||
directOutputSet = newDirect
|
||||
continue
|
||||
}
|
||||
|
||||
newID, newDirect := h.handleTerminalInput(data, sessionID, processID, interactiveRunner, wsw, &ptyWriter, directOutputSet)
|
||||
processID = newID
|
||||
directOutputSet = newDirect
|
||||
}
|
||||
}
|
||||
|
||||
// handleResizeCommand processes a resize message in the terminal bridge loop.
|
||||
// Returns the (possibly updated) processID and directOutputSet flag.
|
||||
func (h *TerminalHandler) handleResizeCommand(
|
||||
payload []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
directOutputSet bool,
|
||||
ptyWriter *io.Writer,
|
||||
) (string, bool) {
|
||||
// Set up PTY access BEFORE handling resize if not already done.
|
||||
// This ensures buffered output is sent to the terminal before the resize
|
||||
// triggers a TUI redraw (which would overwrite the scrollback).
|
||||
if !directOutputSet {
|
||||
if h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
|
||||
// Handle resize — may return updated processID if process was restarted.
|
||||
// Falls back to session-level resize when the old process is gone.
|
||||
newProcessID := h.handleResize(payload, sessionID, processID, interactiveRunner)
|
||||
if newProcessID != processID {
|
||||
processID = newProcessID
|
||||
directOutputSet = false
|
||||
*ptyWriter = nil
|
||||
}
|
||||
|
||||
// Retry setupPtyAccess after resize — process may have been lazily started
|
||||
// by the resize (deferred-start), so the PTY is now available.
|
||||
if !directOutputSet {
|
||||
if h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// handleTerminalInput processes regular (non-resize) input in the terminal bridge loop.
|
||||
// Returns the (possibly updated) processID and directOutputSet flag.
|
||||
func (h *TerminalHandler) handleTerminalInput(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
if *ptyWriter != nil {
|
||||
return h.writeToReadyPty(data, sessionID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
return h.writeToUnreadyPty(data, sessionID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
|
||||
// writeToReadyPty writes input to an already-established PTY writer.
|
||||
// If writing fails it resets state and discovers the new process.
|
||||
func (h *TerminalHandler) writeToReadyPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(err))
|
||||
|
||||
// Reset state and discover the new process immediately.
|
||||
// Without this, the bridge would be stuck with the old processID
|
||||
// until a resize message arrives (which may never come after restart).
|
||||
*ptyWriter = nil
|
||||
directOutputSet = false
|
||||
processID = h.discoverProcessBySession(sessionID, processID, interactiveRunner)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
// Detect Enter key (user submitted input) - notify lifecycle manager
|
||||
h.detectInputSubmission(sessionID, data)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// writeToUnreadyPty attempts to set up PTY access and then write input.
|
||||
// If setup fails it tries to discover the current process by session.
|
||||
func (h *TerminalHandler) writeToUnreadyPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) (string, bool) {
|
||||
// PTY not ready — try setup. If it fails, processID may be stale
|
||||
// (process was restarted with a new ID). Discover the current process
|
||||
// by session before giving up.
|
||||
if !h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
newID := h.discoverProcessBySession(sessionID, processID, interactiveRunner)
|
||||
if newID != processID {
|
||||
processID = newID
|
||||
h.setupPtyAccess(sessionID, processID, interactiveRunner, wsw, ptyWriter)
|
||||
}
|
||||
}
|
||||
if *ptyWriter == nil {
|
||||
h.logger.Debug("PTY not ready, dropping input",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Int("bytes", len(data)))
|
||||
return processID, directOutputSet
|
||||
}
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after setup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
return processID, directOutputSet
|
||||
}
|
||||
h.detectInputSubmission(sessionID, data)
|
||||
return processID, directOutputSet
|
||||
}
|
||||
|
||||
// detectInputSubmission checks if the input contains Enter key and notifies
|
||||
// the lifecycle manager to update the execution state to Running.
|
||||
func (h *TerminalHandler) detectInputSubmission(sessionID string, data []byte) {
|
||||
if !bytes.ContainsAny(data, "\n\r") {
|
||||
return
|
||||
}
|
||||
if err := h.lifecycleMgr.MarkPassthroughRunning(sessionID); err != nil {
|
||||
// May fail if session ended or not in passthrough mode.
|
||||
h.logger.Debug("failed to mark passthrough as running",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
|
||||
// setupPtyAccess attempts to get PTY writer and set up direct output.
|
||||
// Returns true if successful.
|
||||
func (h *TerminalHandler) setupPtyAccess(
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
getWriter := func() (io.Writer, error) {
|
||||
return interactiveRunner.GetPtyWriter(processID)
|
||||
}
|
||||
return h.setupPtyAccessCommon(sessionID, processID, interactiveRunner, wsw, ptyWriter, getWriter)
|
||||
}
|
||||
|
||||
// handleResize processes a resize command from the WebSocket.
|
||||
// Uses process-specific resize to avoid targeting the wrong process when multiple
|
||||
// processes exist for the same session (e.g., passthrough + user shell).
|
||||
// If the process was restarted (old process gone), falls back to session-level resize
|
||||
// and returns the new process ID.
|
||||
func (h *TerminalHandler) handleResize(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) (effectiveProcessID string) {
|
||||
effectiveProcessID = processID
|
||||
|
||||
var resize ResizePayload
|
||||
if err := json.Unmarshal(data, &resize); err != nil {
|
||||
h.logger.Warn("failed to parse resize command",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(err))
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
if resize.Cols == 0 || resize.Rows == 0 {
|
||||
h.logger.Warn("invalid resize dimensions",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
// Resize the specific process by ID to avoid targeting the wrong process.
|
||||
// If the process was restarted (old ID gone), fall back to session-level resize.
|
||||
resizeErr := interactiveRunner.ResizeByProcessID(processID, resize.Cols, resize.Rows)
|
||||
if resizeErr == nil {
|
||||
h.logger.Debug("PTY resized",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
}
|
||||
|
||||
if resizeErr != nil {
|
||||
h.logger.Debug("process-specific resize failed, trying session-level fallback",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(resizeErr))
|
||||
effectiveProcessID = h.resizeBySessionFallback(sessionID, processID, resize.Cols, resize.Rows, interactiveRunner)
|
||||
}
|
||||
|
||||
return effectiveProcessID
|
||||
}
|
||||
|
||||
// discoverProcessBySession looks up the current passthrough process for a session.
|
||||
// Returns the new processID if found and different from oldProcessID.
|
||||
// This is used when the bridge has a stale processID after a process restart.
|
||||
func (h *TerminalHandler) discoverProcessBySession(
|
||||
sessionID string,
|
||||
oldProcessID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) string {
|
||||
_, newProcID, _ := interactiveRunner.GetPtyWriterBySession(sessionID)
|
||||
if newProcID != "" && newProcID != oldProcessID {
|
||||
h.logger.Info("discovered new process by session lookup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("old_process_id", oldProcessID),
|
||||
zap.String("new_process_id", newProcID))
|
||||
return newProcID
|
||||
}
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// resizeBySessionFallback attempts a session-level resize when the process-specific
|
||||
// resize failed (e.g., process was restarted). Returns the new processID if found.
|
||||
func (h *TerminalHandler) resizeBySessionFallback(
|
||||
sessionID string,
|
||||
oldProcessID string,
|
||||
cols, rows uint16,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) string {
|
||||
sessionErr := interactiveRunner.ResizeBySession(sessionID, cols, rows)
|
||||
if sessionErr != nil {
|
||||
h.logger.Warn("failed to resize PTY (both process and session)",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", oldProcessID),
|
||||
zap.Uint16("cols", cols),
|
||||
zap.Uint16("rows", rows),
|
||||
zap.Error(sessionErr))
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// Session-level resize succeeded — look up the new process ID
|
||||
_, newProcID, _ := interactiveRunner.GetPtyWriterBySession(sessionID)
|
||||
if newProcID != "" && newProcID != oldProcessID {
|
||||
h.logger.Info("resize fallback succeeded, process ID updated",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("old_process_id", oldProcessID),
|
||||
zap.String("new_process_id", newProcID))
|
||||
return newProcID
|
||||
}
|
||||
|
||||
return oldProcessID
|
||||
}
|
||||
|
||||
// runUserShellBridge handles WebSocket I/O for user shell terminals.
|
||||
func (h *TerminalHandler) runUserShellBridge(
|
||||
conn *gorillaws.Conn,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
) {
|
||||
var ptyWriter io.Writer
|
||||
var directOutputSet bool
|
||||
|
||||
defer func() {
|
||||
// Clean up WebSocket resources but DON'T stop the process
|
||||
// The process should only be stopped via explicit user_shell.stop message from frontend
|
||||
// This allows reconnection after React remounts or temporary disconnects
|
||||
interactiveRunner.ClearUserShellDirectOutput(scopeID, terminalID)
|
||||
|
||||
_ = wsw.Close()
|
||||
_ = conn.Close()
|
||||
|
||||
h.logger.Info("user shell WebSocket disconnected (process still running for potential reconnect)",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID))
|
||||
}()
|
||||
|
||||
// Read from WebSocket and handle input/resize
|
||||
for {
|
||||
messageType, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
if !gorillaws.IsCloseError(err, gorillaws.CloseNormalClosure, gorillaws.CloseGoingAway) {
|
||||
h.logger.Debug("WebSocket read error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
if messageType != gorillaws.BinaryMessage && messageType != gorillaws.TextMessage {
|
||||
continue
|
||||
}
|
||||
|
||||
if len(data) == 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
if isResizeCommand(data) {
|
||||
directOutputSet = h.handleUserShellResizeCommand(
|
||||
data[1:], sessionID, scopeID, terminalID, processID,
|
||||
interactiveRunner, wsw, directOutputSet, &ptyWriter,
|
||||
)
|
||||
continue
|
||||
}
|
||||
|
||||
directOutputSet = h.handleUserShellInput(
|
||||
data, sessionID, scopeID, terminalID, processID,
|
||||
interactiveRunner, wsw, &ptyWriter, directOutputSet,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
// handleUserShellResizeCommand processes a resize message in the user shell bridge loop.
|
||||
// Returns the updated directOutputSet flag.
|
||||
func (h *TerminalHandler) handleUserShellResizeCommand(
|
||||
payload []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
directOutputSet bool,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
// Set up PTY access BEFORE handling resize if not already done
|
||||
if !directOutputSet {
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
}
|
||||
}
|
||||
h.handleUserShellResize(payload, sessionID, scopeID, terminalID, interactiveRunner)
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// handleUserShellInput processes regular (non-resize) input in the user shell bridge loop.
|
||||
// Returns the updated directOutputSet flag.
|
||||
func (h *TerminalHandler) handleUserShellInput(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if *ptyWriter != nil {
|
||||
return h.writeToReadyUserShellPty(data, sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
return h.writeToUnreadyUserShellPty(data, sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter, directOutputSet)
|
||||
}
|
||||
|
||||
// writeToReadyUserShellPty writes input to an established user shell PTY.
|
||||
// On failure it attempts to re-establish PTY access and retries the write.
|
||||
func (h *TerminalHandler) writeToReadyUserShellPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
// Try to re-establish PTY access
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after reconnect",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
}
|
||||
}
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// writeToUnreadyUserShellPty sets up PTY access and writes input for a user shell.
|
||||
// Drops the input if setup fails.
|
||||
func (h *TerminalHandler) writeToUnreadyUserShellPty(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
directOutputSet bool,
|
||||
) bool {
|
||||
if h.setupUserShellPtyAccess(sessionID, scopeID, terminalID, processID, interactiveRunner, wsw, ptyWriter) {
|
||||
directOutputSet = true
|
||||
if _, err := (*ptyWriter).Write(data); err != nil {
|
||||
h.logger.Debug("PTY write error after setup",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
}
|
||||
} else {
|
||||
h.logger.Debug("PTY not ready, dropping input",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Int("bytes", len(data)))
|
||||
}
|
||||
return directOutputSet
|
||||
}
|
||||
|
||||
// setupUserShellPtyAccess sets up PTY access for a user shell terminal.
|
||||
func (h *TerminalHandler) setupUserShellPtyAccess(
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
) bool {
|
||||
getWriter := func() (io.Writer, error) {
|
||||
writer, _, err := interactiveRunner.GetUserShellPtyWriter(scopeID, terminalID)
|
||||
return writer, err
|
||||
}
|
||||
return h.setupPtyAccessCommon(sessionID, processID, interactiveRunner, wsw, ptyWriter, getWriter)
|
||||
}
|
||||
|
||||
// replayBufferedOutput flushes any buffered PTY output for a process to the WebSocket
|
||||
// before direct output is wired up. This is what gives a reconnecting terminal its
|
||||
// scrollback history. Terminal query response sequences are stripped first so they
|
||||
// don't render as visible garbage on replay.
|
||||
func (h *TerminalHandler) replayBufferedOutput(
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
) {
|
||||
chunks, ok := interactiveRunner.GetBuffer(processID)
|
||||
if !ok || len(chunks) == 0 {
|
||||
return
|
||||
}
|
||||
var combined bytes.Buffer
|
||||
for _, chunk := range chunks {
|
||||
combined.WriteString(chunk.Data)
|
||||
}
|
||||
if combined.Len() == 0 {
|
||||
return
|
||||
}
|
||||
cleaned := stripTerminalResponses(combined.Bytes())
|
||||
h.logger.Debug("sending buffered output to terminal",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Int("chunks", len(chunks)),
|
||||
zap.Int("total_bytes", combined.Len()),
|
||||
zap.Int("cleaned_bytes", len(cleaned)))
|
||||
if len(cleaned) == 0 {
|
||||
return
|
||||
}
|
||||
if _, writeErr := wsw.Write(cleaned); writeErr != nil {
|
||||
h.logger.Warn("failed to send buffered output",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Error(writeErr))
|
||||
}
|
||||
}
|
||||
|
||||
// setupPtyAccessCommon is the shared implementation for setupPtyAccess and setupUserShellPtyAccess.
|
||||
// It gets a PTY writer via getWriter, sends buffered output, and sets up direct output.
|
||||
func (h *TerminalHandler) setupPtyAccessCommon(
|
||||
sessionID string,
|
||||
processID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
wsw *wsWriter,
|
||||
ptyWriter *io.Writer,
|
||||
getWriter func() (io.Writer, error),
|
||||
) bool {
|
||||
writer, err := getWriter()
|
||||
if err != nil {
|
||||
h.logger.Debug("PTY writer not ready yet",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID))
|
||||
return false
|
||||
}
|
||||
|
||||
// Send buffered output to the terminal before switching to direct mode.
|
||||
// This provides scrollback history when reconnecting to an existing session.
|
||||
h.replayBufferedOutput(sessionID, processID, interactiveRunner, wsw)
|
||||
|
||||
// Set up direct output
|
||||
if err := interactiveRunner.SetDirectOutput(processID, wsw); err != nil {
|
||||
h.logger.Error("failed to set direct output",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID),
|
||||
zap.Error(err))
|
||||
return false
|
||||
}
|
||||
|
||||
*ptyWriter = writer
|
||||
|
||||
h.logger.Info("PTY access established",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("process_id", processID))
|
||||
|
||||
return true
|
||||
}
|
||||
|
||||
// handleUserShellResize processes resize commands for user shell terminals.
|
||||
func (h *TerminalHandler) handleUserShellResize(
|
||||
data []byte,
|
||||
sessionID string,
|
||||
scopeID string,
|
||||
terminalID string,
|
||||
interactiveRunner *process.InteractiveRunner,
|
||||
) {
|
||||
var resize ResizePayload
|
||||
if err := json.Unmarshal(data, &resize); err != nil {
|
||||
h.logger.Warn("failed to parse resize command",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Error(err))
|
||||
return
|
||||
}
|
||||
|
||||
if resize.Cols == 0 || resize.Rows == 0 {
|
||||
h.logger.Warn("invalid resize dimensions",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
return
|
||||
}
|
||||
|
||||
if err := interactiveRunner.ResizeUserShell(scopeID, terminalID, resize.Cols, resize.Rows); err != nil {
|
||||
h.logger.Warn("failed to resize user shell PTY",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows),
|
||||
zap.Error(err))
|
||||
} else {
|
||||
h.logger.Debug("user shell PTY resized",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("terminal_id", terminalID),
|
||||
zap.Uint16("cols", resize.Cols),
|
||||
zap.Uint16("rows", resize.Rows))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,133 @@
|
||||
package websocket
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/kandev/kandev/internal/agent/runtime/lifecycle"
|
||||
)
|
||||
|
||||
const (
|
||||
passthroughReadyTimeout = 30 * time.Second
|
||||
remoteReadyProbeTimeout = 500 * time.Millisecond
|
||||
pollInterval = 500 * time.Millisecond
|
||||
)
|
||||
|
||||
func (h *TerminalHandler) shouldUseWorkspaceShell(ctx context.Context, sessionID string) bool {
|
||||
// Delegates to lifecycle manager which checks both in-memory execution
|
||||
// and database records for containerized/remote executors.
|
||||
return h.lifecycleMgr.ShouldUseContainerShell(ctx, sessionID)
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) waitForRemoteExecutionReadyWithTimeout(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
timeout time.Duration,
|
||||
) (*lifecycle.AgentExecution, bool) {
|
||||
deadline := time.Now().Add(timeout)
|
||||
|
||||
h.logger.Info("waiting for remote execution to become ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", timeout))
|
||||
|
||||
for {
|
||||
if execution, ok := h.remoteExecutionWithClientReady(ctx, sessionID, deadline); ok {
|
||||
return execution, true
|
||||
}
|
||||
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
h.logger.Warn("timed out waiting for remote execution readiness",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", timeout))
|
||||
return nil, false
|
||||
}
|
||||
|
||||
wait := pollInterval
|
||||
if remaining < wait {
|
||||
wait = remaining
|
||||
}
|
||||
timer := time.NewTimer(wait)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
timer.Stop()
|
||||
return nil, false
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) remoteExecutionWithClientReady(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
deadline time.Time,
|
||||
) (*lifecycle.AgentExecution, bool) {
|
||||
execution, exists := h.lifecycleMgr.GetExecutionBySessionID(sessionID)
|
||||
if !exists || execution.GetAgentCtlClient() == nil {
|
||||
return nil, false
|
||||
}
|
||||
remaining := time.Until(deadline)
|
||||
if remaining <= 0 {
|
||||
return nil, false
|
||||
}
|
||||
|
||||
probeTimeout := remoteReadyProbeTimeout
|
||||
if remaining < probeTimeout {
|
||||
probeTimeout = remaining
|
||||
}
|
||||
probeCtx, cancel := context.WithTimeout(ctx, probeTimeout)
|
||||
defer cancel()
|
||||
if err := execution.GetAgentCtlClient().Health(probeCtx); err != nil {
|
||||
h.logger.Debug("remote execution agentctl client not ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID),
|
||||
zap.Error(err))
|
||||
return nil, false
|
||||
}
|
||||
h.logger.Info("remote execution became ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID))
|
||||
return execution, true
|
||||
}
|
||||
|
||||
func (h *TerminalHandler) ensurePassthroughExecutionReady(ctx context.Context, sessionID string) (*lifecycle.AgentExecution, error) {
|
||||
execution, err := h.lifecycleMgr.EnsurePassthroughExecution(ctx, sessionID)
|
||||
if err == nil {
|
||||
return execution, nil
|
||||
}
|
||||
if !errors.Is(err, lifecycle.ErrSessionWorkspaceNotReady) {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
h.logger.Info("waiting for passthrough workspace to become ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.Duration("timeout", passthroughReadyTimeout))
|
||||
|
||||
timeoutCtx, cancel := context.WithTimeout(ctx, passthroughReadyTimeout)
|
||||
defer cancel()
|
||||
|
||||
ticker := time.NewTicker(pollInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-timeoutCtx.Done():
|
||||
return nil, fmt.Errorf("%w: timed out after %s", lifecycle.ErrSessionWorkspaceNotReady, passthroughReadyTimeout)
|
||||
case <-ticker.C:
|
||||
execution, err := h.lifecycleMgr.EnsurePassthroughExecution(timeoutCtx, sessionID)
|
||||
if err == nil {
|
||||
h.logger.Info("passthrough workspace became ready",
|
||||
zap.String("session_id", sessionID),
|
||||
zap.String("execution_id", execution.ID))
|
||||
return execution, nil
|
||||
}
|
||||
if !errors.Is(err, lifecycle.ErrSessionWorkspaceNotReady) {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user