feat: improve error handling for WebSocket message sending and enhance API documentation

Author: pr0head

docs/src/content/docs/reference/admin-api.md

@@ -1013,7 +1013,7 @@ Send a text or binary message to a specific active WebSocket connection.
 | Field | Type | Description |
 |-------|------|-------------|
 | `type` | string | Message type: `"text"` (default) or `"binary"` |
-| `data` | string | Message payload |
+| `data` | string | Message payload. For `"text"`, a plain UTF-8 string. For `"binary"`, a **base64-encoded** string — the server decodes it before writing raw bytes to the WebSocket. |
 
 **Response:**
 

pkg/admin/engineclient/types.go

@@ -30,32 +30,31 @@ type (
 	ProtocolStatus   = types.ProtocolStatus
 	MockListResponse = types.MockListResponse
 	// RequestFilter is kept for backward compatibility; use requestlog.Filter directly.
-	RequestFilter                   = types.RequestLogFilter
-	RequestLogEntry                 = types.RequestLogEntry
-	RequestListResponse             = types.RequestListResponse
-	ErrorResponse                   = types.ErrorResponse
-	ChaosConfig                     = types.ChaosConfig
-	LatencyConfig                   = types.LatencyConfig
-	ErrorRateConfig                 = types.ErrorRateConfig
-	BandwidthConfig                 = types.BandwidthConfig
-	ChaosRuleConfig                 = types.ChaosRuleConfig
-	ChaosFaultConfig                = types.ChaosFaultConfig
-	ChaosStats                      = types.ChaosStats
-	StatefulResource                = types.StatefulResource
-	StatefulItemsResponse           = types.StatefulItemsResponse
-	StateOverview                   = types.StateOverview
-	ProtocolHandler                 = types.ProtocolHandler
-	SSEConnection                   = types.SSEConnection
-	SSEStats                        = types.SSEStats
-	WebSocketConnection             = types.WebSocketConnection
-	WebSocketConnectionListResponse = types.WebSocketConnectionListResponse
-	WebSocketStats                  = types.WebSocketStats
-	CustomOperationInfo             = types.CustomOperationInfo
-	CustomOperationDetail           = types.CustomOperationDetail
-	CustomOperationStep             = types.CustomOperationStep
-	StatefulFaultStats              = types.StatefulFaultStats
-	CircuitBreakerStatus            = types.CircuitBreakerStatus
-	RetryAfterStatus                = types.RetryAfterStatus
-	ProgressiveDegradationStatus    = types.ProgressiveDegradationStatus
-	ResetStateResponse              = types.ResetStateResponse
+	RequestFilter                = types.RequestLogFilter
+	RequestLogEntry              = types.RequestLogEntry
+	RequestListResponse          = types.RequestListResponse
+	ErrorResponse                = types.ErrorResponse
+	ChaosConfig                  = types.ChaosConfig
+	LatencyConfig                = types.LatencyConfig
+	ErrorRateConfig              = types.ErrorRateConfig
+	BandwidthConfig              = types.BandwidthConfig
+	ChaosRuleConfig              = types.ChaosRuleConfig
+	ChaosFaultConfig             = types.ChaosFaultConfig
+	ChaosStats                   = types.ChaosStats
+	StatefulResource             = types.StatefulResource
+	StatefulItemsResponse        = types.StatefulItemsResponse
+	StateOverview                = types.StateOverview
+	ProtocolHandler              = types.ProtocolHandler
+	SSEConnection                = types.SSEConnection
+	SSEStats                     = types.SSEStats
+	WebSocketConnection          = types.WebSocketConnection
+	WebSocketStats               = types.WebSocketStats
+	CustomOperationInfo          = types.CustomOperationInfo
+	CustomOperationDetail        = types.CustomOperationDetail
+	CustomOperationStep          = types.CustomOperationStep
+	StatefulFaultStats           = types.StatefulFaultStats
+	CircuitBreakerStatus         = types.CircuitBreakerStatus
+	RetryAfterStatus             = types.RetryAfterStatus
+	ProgressiveDegradationStatus = types.ProgressiveDegradationStatus
+	ResetStateResponse           = types.ResetStateResponse
 )

pkg/admin/sse_handlers.go

@@ -150,8 +150,11 @@ func (a *API) handleCloseSSEConnection(w http.ResponseWriter, r *http.Request) {
 // handleGetSSEStats handles GET /sse/stats.
 func (a *API) handleGetSSEStats(w http.ResponseWriter, r *http.Request) {
 	engine := a.localEngine.Load()
-	provider := newSSEStatsProvider(engine)
-	a.handleGetStats(w, r, provider)
+	if engine == nil {
+		writeJSON(w, http.StatusOK, sse.ConnectionStats{ConnectionsByMock: make(map[string]int)})
+		return
+	}
+	a.handleGetStats(w, r, newSSEStatsProvider(engine))
 }
 
 // handleListMockSSEConnections handles GET /mocks/{id}/sse/connections.

pkg/admin/stat_helper.go

@@ -20,16 +20,11 @@ type statsProvider interface {
 }
 
 // handleGetStats is a generic handler for retrieving statistics.
-// It eliminates code duplication between SSE and WebSocket stats handlers.
+// The caller must guard against a nil engine and handle the empty-stats
+// response before constructing the provider and invoking this function.
 func (a *API) handleGetStats(w http.ResponseWriter, r *http.Request, provider statsProvider) {
 	ctx := r.Context()
 
-	engine := a.localEngine.Load()
-	if engine == nil {
-		writeJSON(w, http.StatusOK, provider.GetEmptyStats())
-		return
-	}
-
 	stats, err := provider.GetStats(ctx)
 	if err != nil {
 		a.logger().Error("failed to get stats", "error", err)

pkg/admin/websocket_handlers.go

@@ -134,8 +134,11 @@ func (a *API) handleCloseWebSocketConnection(w http.ResponseWriter, r *http.Requ
 // handleGetWebSocketStats handles GET /websocket/stats.
 func (a *API) handleGetWebSocketStats(w http.ResponseWriter, r *http.Request) {
 	engine := a.localEngine.Load()
-	provider := newWSStatsProvider(engine)
-	a.handleGetStats(w, r, provider)
+	if engine == nil {
+		writeJSON(w, http.StatusOK, engineclient.WebSocketStats{ConnectionsByMock: make(map[string]int)})
+		return
+	}
+	a.handleGetStats(w, r, newWSStatsProvider(engine))
 }
 
 func mapWebSocketEngineError(err error, log *slog.Logger, operation string) (int, string, string) {

pkg/engine/api/handlers.go

@@ -15,6 +15,7 @@ import (
 	"github.com/getmockd/mockd/pkg/httputil"
 	"github.com/getmockd/mockd/pkg/requestlog"
 	"github.com/getmockd/mockd/pkg/stateful"
+	"github.com/getmockd/mockd/pkg/websocket"
 )
 
 func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
@@ -856,7 +857,11 @@ func (s *Server) handleSendToWebSocketConnection(w http.ResponseWriter, r *http.
 	}
 
 	if err := s.engine.SendToWebSocketConnection(id, req.Type, payload); err != nil {
-		writeError(w, http.StatusNotFound, "not_found", "WebSocket connection not found")
+		if errors.Is(err, websocket.ErrConnectionNotFound) || errors.Is(err, websocket.ErrConnectionClosed) {
+			writeError(w, http.StatusNotFound, "not_found", "WebSocket connection not found or closed")
+		} else {
+			writeError(w, http.StatusInternalServerError, "send_error", "Failed to send message to WebSocket connection")
+		}
 		return
 	}
 

pkg/engine/api/handlers_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/getmockd/mockd/pkg/mock"
 	"github.com/getmockd/mockd/pkg/requestlog"
 	"github.com/getmockd/mockd/pkg/stateful"
+	"github.com/getmockd/mockd/pkg/websocket"
 )
 
 // mockEngine is a test double for EngineController.
@@ -49,13 +50,17 @@ type mockEngine struct {
 	resetStateErr         error
 	listStatefulItemsErr  error
 	getStatefulItemErr    error
+	// wsSendErr allows injecting a specific error for a connection ID in
+	// SendToWebSocketConnection. Keyed by connection ID.
+	wsSendErr map[string]error
 }
 
 func newMockEngine() *mockEngine {
 	return &mockEngine{
 		mocks:       make(map[string]*config.MockConfiguration),
 		requestLogs: make(map[string]*requestlog.Entry),
 		customOps:   make(map[string]*CustomOperationDetail),
+		wsSendErr:   make(map[string]error),
 		running:     true,
 		uptime:      100,
 		protocols: map[string]ProtocolStatusInfo{
@@ -360,12 +365,16 @@ func (m *mockEngine) GetWebSocketStats() *WebSocketStats {
 }
 
 func (m *mockEngine) SendToWebSocketConnection(id string, msgType string, data []byte) error {
+	// Allow per-connection error injection for testing specific error paths.
+	if err, ok := m.wsSendErr[id]; ok {
+		return err
+	}
 	for _, c := range m.wsConnections {
 		if c.ID == id {
 			return nil
 		}
 	}
-	return errors.New("connection not found")
+	return websocket.ErrConnectionNotFound
 }
 
 func (m *mockEngine) GetConfig() *ConfigResponse {
@@ -2406,3 +2415,169 @@ func TestHandleGetWebSocketStats(t *testing.T) {
 		assert.Equal(t, 3, stats.ConnectionsByMock["mock-1"])
 	})
 }
+
+// TestHandleSendToWebSocketConnection covers every branch of the send handler,
+// including the new 404/500 distinction introduced in the review fix.
+func TestHandleSendToWebSocketConnection(t *testing.T) {
+	t.Run("returns 400 when connection id is missing", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections//send",
+			strings.NewReader(`{"type":"text","data":"hello"}`))
+		// PathValue "id" intentionally not set → empty string
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusBadRequest, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "missing_id", resp["error"])
+	})
+
+	t.Run("returns 400 for invalid JSON body", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{invalid`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusBadRequest, rec.Code)
+	})
+
+	t.Run("returns 400 for invalid base64 binary payload", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		engine.wsConnections = []*WebSocketConnection{{ID: "ws-1"}}
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"type":"binary","data":"not-valid-base64!!!"}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusBadRequest, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "invalid_base64", resp["error"])
+	})
+
+	t.Run("returns 200 for text message to existing connection", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		engine.wsConnections = []*WebSocketConnection{{ID: "ws-1"}}
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"type":"text","data":"hello"}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "Message sent", resp["message"])
+		assert.Equal(t, "ws-1", resp["connection"])
+		assert.Equal(t, "text", resp["type"])
+	})
+
+	t.Run("returns 200 for valid base64 binary message", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		engine.wsConnections = []*WebSocketConnection{{ID: "ws-1"}}
+
+		// base64("hello") = "aGVsbG8="
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"type":"binary","data":"aGVsbG8="}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "binary", resp["type"])
+	})
+
+	t.Run("defaults type to text when omitted", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		engine.wsConnections = []*WebSocketConnection{{ID: "ws-1"}}
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"data":"hello"}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusOK, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "text", resp["type"])
+	})
+
+	t.Run("returns 404 when connection is not found", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		// No connections registered → ErrConnectionNotFound
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/missing/send",
+			strings.NewReader(`{"type":"text","data":"hello"}`))
+		req.SetPathValue("id", "missing")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusNotFound, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "not_found", resp["error"])
+	})
+
+	t.Run("returns 404 when connection is already closed", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		// Inject ErrConnectionClosed for ws-1 (connection exists but is closed)
+		engine.wsSendErr["ws-1"] = websocket.ErrConnectionClosed
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"type":"text","data":"hello"}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusNotFound, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "not_found", resp["error"])
+	})
+
+	t.Run("returns 500 for unexpected write error", func(t *testing.T) {
+		engine := newMockEngine()
+		server := newTestServer(engine)
+		// Inject a generic I/O error (e.g., broken pipe)
+		engine.wsSendErr["ws-1"] = errors.New("write: broken pipe")
+
+		req := httptest.NewRequest(http.MethodPost, "/websocket/connections/ws-1/send",
+			strings.NewReader(`{"type":"text","data":"hello"}`))
+		req.SetPathValue("id", "ws-1")
+		rec := httptest.NewRecorder()
+
+		server.handleSendToWebSocketConnection(rec, req)
+
+		assert.Equal(t, http.StatusInternalServerError, rec.Code)
+		var resp map[string]interface{}
+		require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp))
+		assert.Equal(t, "send_error", resp["error"])
+	})
+}