refactor: unify agent communication with Transport interface

- Introduce `Transport` interface to abstract WebSocket and SSH
communication
- Add generic `Data` field to `AgentResponse` for streamlined future
endpoints
- Maintain backward compatibility with legacy hubs and agents using
typed fields
- Unify fetch operations (SMART, systemd, containers) under a single
`request` method
- Improve `RequestManager` with deadline awareness and legacy response
support
- Refactor agent response routing into dedicated `agent/response.go`
- Update version to 0.18.0-beta.2

internal/hub/ws/handlers.go

@@ -6,14 +6,12 @@ import (
 
 	"github.com/fxamacker/cbor/v2"
 	"github.com/henrygd/beszel/internal/common"
-	"github.com/henrygd/beszel/internal/entities/smart"
-	"github.com/henrygd/beszel/internal/entities/system"
-	"github.com/henrygd/beszel/internal/entities/systemd"
 	"github.com/lxzan/gws"
 	"golang.org/x/crypto/ssh"
 )
 
-// ResponseHandler defines interface for handling agent responses
+// ResponseHandler defines interface for handling agent responses.
+// This is used by handleAgentRequest for legacy response handling.
 type ResponseHandler interface {
 	Handle(agentResponse common.AgentResponse) error
 	HandleLegacy(rawData []byte) error
@@ -27,167 +25,7 @@ func (h *BaseHandler) HandleLegacy(rawData []byte) error {
 }
 
 ////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-
-// systemDataHandler implements ResponseHandler for system data requests
-type systemDataHandler struct {
-	data *system.CombinedData
-}
-
-func (h *systemDataHandler) HandleLegacy(rawData []byte) error {
-	return cbor.Unmarshal(rawData, h.data)
-}
-
-func (h *systemDataHandler) Handle(agentResponse common.AgentResponse) error {
-	if agentResponse.SystemData != nil {
-		*h.data = *agentResponse.SystemData
-	}
-	return nil
-}
-
-// RequestSystemData requests system metrics from the agent and unmarshals the response.
-func (ws *WsConn) RequestSystemData(ctx context.Context, data *system.CombinedData, options common.DataRequestOptions) error {
-	if !ws.IsConnected() {
-		return gws.ErrConnClosed
-	}
-
-	req, err := ws.requestManager.SendRequest(ctx, common.GetData, options)
-	if err != nil {
-		return err
-	}
-
-	handler := &systemDataHandler{data: data}
-	return ws.handleAgentRequest(req, handler)
-}
-
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-
-// stringResponseHandler is a generic handler for string responses from agents
-type stringResponseHandler struct {
-	BaseHandler
-	value    string
-	errorMsg string
-}
-
-func (h *stringResponseHandler) Handle(agentResponse common.AgentResponse) error {
-	if agentResponse.String == nil {
-		return errors.New(h.errorMsg)
-	}
-	h.value = *agentResponse.String
-	return nil
-}
-
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-
-// requestContainerStringViaWS is a generic function to request container-related strings via WebSocket
-func (ws *WsConn) requestContainerStringViaWS(ctx context.Context, action common.WebSocketAction, requestData any, errorMsg string) (string, error) {
-	if !ws.IsConnected() {
-		return "", gws.ErrConnClosed
-	}
-
-	req, err := ws.requestManager.SendRequest(ctx, action, requestData)
-	if err != nil {
-		return "", err
-	}
-
-	handler := &stringResponseHandler{errorMsg: errorMsg}
-	if err := ws.handleAgentRequest(req, handler); err != nil {
-		return "", err
-	}
-
-	return handler.value, nil
-}
-
-// RequestContainerLogs requests logs for a specific container via WebSocket.
-func (ws *WsConn) RequestContainerLogs(ctx context.Context, containerID string) (string, error) {
-	return ws.requestContainerStringViaWS(ctx, common.GetContainerLogs, common.ContainerLogsRequest{ContainerID: containerID}, "no logs in response")
-}
-
-// RequestContainerInfo requests information about a specific container via WebSocket.
-func (ws *WsConn) RequestContainerInfo(ctx context.Context, containerID string) (string, error) {
-	return ws.requestContainerStringViaWS(ctx, common.GetContainerInfo, common.ContainerInfoRequest{ContainerID: containerID}, "no info in response")
-}
-
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-
-// RequestSystemdInfo requests detailed information about a systemd service via WebSocket.
-func (ws *WsConn) RequestSystemdInfo(ctx context.Context, serviceName string) (systemd.ServiceDetails, error) {
-	if !ws.IsConnected() {
-		return nil, gws.ErrConnClosed
-	}
-
-	req, err := ws.requestManager.SendRequest(ctx, common.GetSystemdInfo, common.SystemdInfoRequest{ServiceName: serviceName})
-	if err != nil {
-		return nil, err
-	}
-
-	var result systemd.ServiceDetails
-	handler := &systemdInfoHandler{result: &result}
-	if err := ws.handleAgentRequest(req, handler); err != nil {
-		return nil, err
-	}
-
-	return result, nil
-}
-
-// systemdInfoHandler parses ServiceDetails from AgentResponse
-type systemdInfoHandler struct {
-	BaseHandler
-	result *systemd.ServiceDetails
-}
-
-func (h *systemdInfoHandler) Handle(agentResponse common.AgentResponse) error {
-	if agentResponse.ServiceInfo == nil {
-		return errors.New("no systemd info in response")
-	}
-	*h.result = agentResponse.ServiceInfo
-	return nil
-}
-
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
-
-// RequestSmartData requests SMART data via WebSocket.
-func (ws *WsConn) RequestSmartData(ctx context.Context) (map[string]smart.SmartData, error) {
-	if !ws.IsConnected() {
-		return nil, gws.ErrConnClosed
-	}
-	req, err := ws.requestManager.SendRequest(ctx, common.GetSmartData, nil)
-	if err != nil {
-		return nil, err
-	}
-	var result map[string]smart.SmartData
-	handler := ResponseHandler(&smartDataHandler{result: &result})
-	if err := ws.handleAgentRequest(req, handler); err != nil {
-		return nil, err
-	}
-	return result, nil
-}
-
-// smartDataHandler parses SMART data map from AgentResponse
-type smartDataHandler struct {
-	BaseHandler
-	result *map[string]smart.SmartData
-}
-
-func (h *smartDataHandler) Handle(agentResponse common.AgentResponse) error {
-	if agentResponse.SmartData == nil {
-		return errors.New("no SMART data in response")
-	}
-	*h.result = agentResponse.SmartData
-	return nil
-}
-
-////////////////////////////////////////////////////////////////////////////
-////////////////////////////////////////////////////////////////////////////
+// Fingerprint handling (used for WebSocket authentication)
 ////////////////////////////////////////////////////////////////////////////
 
 // fingerprintHandler implements ResponseHandler for fingerprint requests

internal/hub/ws/handlers_test.go

@@ -1,75 +0,0 @@
-//go:build testing
-
-package ws
-
-import (
-	"testing"
-
-	"github.com/henrygd/beszel/internal/common"
-	"github.com/henrygd/beszel/internal/entities/systemd"
-	"github.com/stretchr/testify/assert"
-)
-
-func TestSystemdInfoHandlerSuccess(t *testing.T) {
-	handler := &systemdInfoHandler{
-		result: &systemd.ServiceDetails{},
-	}
-
-	// Test successful handling with valid ServiceInfo
-	testDetails := systemd.ServiceDetails{
-		"Id":             "nginx.service",
-		"ActiveState":    "active",
-		"SubState":       "running",
-		"Description":    "A high performance web server",
-		"ExecMainPID":    1234,
-		"MemoryCurrent":  1024000,
-	}
-
-	response := common.AgentResponse{
-		ServiceInfo: testDetails,
-	}
-
-	err := handler.Handle(response)
-	assert.NoError(t, err)
-	assert.Equal(t, testDetails, *handler.result)
-}
-
-func TestSystemdInfoHandlerError(t *testing.T) {
-	handler := &systemdInfoHandler{
-		result: &systemd.ServiceDetails{},
-	}
-
-	// Test error handling when ServiceInfo is nil
-	response := common.AgentResponse{
-		ServiceInfo: nil,
-		Error:       "service not found",
-	}
-
-	err := handler.Handle(response)
-	assert.Error(t, err)
-	assert.Equal(t, "no systemd info in response", err.Error())
-}
-
-func TestSystemdInfoHandlerEmptyResponse(t *testing.T) {
-	handler := &systemdInfoHandler{
-		result: &systemd.ServiceDetails{},
-	}
-
-	// Test with completely empty response
-	response := common.AgentResponse{}
-
-	err := handler.Handle(response)
-	assert.Error(t, err)
-	assert.Equal(t, "no systemd info in response", err.Error())
-}
-
-func TestSystemdInfoHandlerLegacyNotSupported(t *testing.T) {
-	handler := &systemdInfoHandler{
-		result: &systemd.ServiceDetails{},
-	}
-
-	// Test that legacy format is not supported
-	err := handler.HandleLegacy([]byte("some data"))
-	assert.Error(t, err)
-	assert.Equal(t, "legacy format not supported", err.Error())
-}

internal/hub/ws/request_manager.go

@@ -45,7 +45,15 @@ func NewRequestManager(conn *gws.Conn) *RequestManager {
 func (rm *RequestManager) SendRequest(ctx context.Context, action common.WebSocketAction, data any) (*PendingRequest, error) {
 	reqID := RequestID(rm.nextID.Add(1))
 
-	reqCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	// Respect any caller-provided deadline. If none is set, apply a reasonable default
+	// so pending requests don't live forever if the agent never responds.
+	reqCtx := ctx
+	var cancel context.CancelFunc
+	if _, hasDeadline := ctx.Deadline(); hasDeadline {
+		reqCtx, cancel = context.WithCancel(ctx)
+	} else {
+		reqCtx, cancel = context.WithTimeout(ctx, 5*time.Second)
+	}
 
 	req := &PendingRequest{
 		ID:         reqID,
@@ -100,6 +108,11 @@ func (rm *RequestManager) handleResponse(message *gws.Message) {
 		return
 	}
 
+	if response.Id == nil {
+		rm.routeLegacyResponse(message)
+		return
+	}
+
 	reqID := RequestID(*response.Id)
 
 	rm.RLock()

internal/hub/ws/ws.go

@@ -1,6 +1,7 @@
 package ws
 
 import (
+	"context"
 	"errors"
 	"time"
 	"weak"
@@ -161,3 +162,14 @@ func (ws *WsConn) handleAgentRequest(req *PendingRequest, handler ResponseHandle
 func (ws *WsConn) IsConnected() bool {
 	return ws.conn != nil
 }
+
+// AgentVersion returns the connected agent's version (as reported during handshake).
+func (ws *WsConn) AgentVersion() semver.Version {
+	return ws.agentVersion
+}
+
+// SendRequest sends a request to the agent and returns a pending request handle.
+// This is used by the transport layer to send requests.
+func (ws *WsConn) SendRequest(ctx context.Context, action common.WebSocketAction, data any) (*PendingRequest, error) {
+	return ws.requestManager.SendRequest(ctx, action, data)
+}

internal/hub/ws/ws_test.go

@@ -184,14 +184,18 @@ func TestCommonActions(t *testing.T) {
 	assert.Equal(t, common.WebSocketAction(2), common.GetContainerLogs, "GetLogs should be action 2")
 }
 
-func TestLogsHandler(t *testing.T) {
-	h := &stringResponseHandler{errorMsg: "no logs in response"}
+func TestFingerprintHandler(t *testing.T) {
+	var result common.FingerprintResponse
+	h := &fingerprintHandler{result: &result}
 
-	logValue := "test logs"
-	resp := common.AgentResponse{String: &logValue}
+	resp := common.AgentResponse{Fingerprint: &common.FingerprintResponse{
+		Fingerprint: "test-fingerprint",
+		Hostname:    "test-host",
+	}}
 	err := h.Handle(resp)
 	assert.NoError(t, err)
-	assert.Equal(t, logValue, h.value)
+	assert.Equal(t, "test-fingerprint", result.Fingerprint)
+	assert.Equal(t, "test-host", result.Hostname)
 }
 
 // TestHandler tests that we can create a Handler