agent: fix container logs decoding for raw streams (#1535)

* NVML: only read memory usage if utilization > 0% to allow rtd3, #1522

* logging: /x/exp/slog -> log/slog everywhere, fixes log instance inconsistencies

agent/client.go

@@ -15,9 +15,6 @@ import (
 
 	"github.com/henrygd/beszel"
 	"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/fxamacker/cbor/v2"
 	"github.com/lxzan/gws"
@@ -259,40 +256,16 @@ func (client *WebSocketClient) sendMessage(data any) error {
 	return err
 }
 
-// sendResponse sends a response with optional request ID for the new protocol
+// sendResponse sends a response with optional request ID.
+// For ID-based requests, we must populate legacy typed fields for backward
+// compatibility with older hubs (<= 0.17) that don't read the generic Data field.
 func (client *WebSocketClient) sendResponse(data any, requestID *uint32) error {
 	if requestID != nil {
-		// New format with ID - use typed fields
-		response := common.AgentResponse{
-			Id: requestID,
-		}
-
-		// Set the appropriate typed field based on data type
-		switch v := data.(type) {
-		case *system.CombinedData:
-			response.SystemData = v
-		case *common.FingerprintResponse:
-			response.Fingerprint = v
-		case string:
-			response.String = &v
-		case map[string]smart.SmartData:
-			response.SmartData = v
-		case systemd.ServiceDetails:
-			response.ServiceInfo = v
-		// case []byte:
-		// 	response.RawBytes = v
-		// case string:
-		// 	response.RawBytes = []byte(v)
-		default:
-			// For any other type, convert to error
-			response.Error = fmt.Sprintf("unsupported response type: %T", data)
-		}
-
+		response := newAgentResponse(data, requestID)
 		return client.sendMessage(response)
-	} else {
-		// Legacy format - send data directly
-		return client.sendMessage(data)
 	}
+	// Legacy format - send data directly
+	return client.sendMessage(data)
 }
 
 // getUserAgent returns one of two User-Agent strings based on current time.

agent/docker.go

@@ -694,7 +694,8 @@ func (dm *dockerManager) getLogs(ctx context.Context, containerID string) (strin
 	}
 
 	var builder strings.Builder
-	if err := decodeDockerLogStream(resp.Body, &builder); err != nil {
+	multiplexed := resp.Header.Get("Content-Type") == "application/vnd.docker.multiplexed-stream"
+	if err := decodeDockerLogStream(resp.Body, &builder, multiplexed); err != nil {
 		return "", err
 	}
 
@@ -706,7 +707,11 @@ func (dm *dockerManager) getLogs(ctx context.Context, containerID string) (strin
 	return logs, nil
 }
 
-func decodeDockerLogStream(reader io.Reader, builder *strings.Builder) error {
+func decodeDockerLogStream(reader io.Reader, builder *strings.Builder, multiplexed bool) error {
+	if !multiplexed {
+		_, err := io.Copy(builder, io.LimitReader(reader, maxTotalLogSize))
+		return err
+	}
 	const headerSize = 8
 	var header [headerSize]byte
 	totalBytesRead := 0

agent/docker_test.go

@@ -950,6 +950,7 @@ func TestDecodeDockerLogStream(t *testing.T) {
 		input       []byte
 		expected    string
 		expectError bool
+		multiplexed bool
 	}{
 		{
 			name: "simple log entry",
@@ -960,6 +961,7 @@ func TestDecodeDockerLogStream(t *testing.T) {
 			},
 			expected:    "Hello World",
 			expectError: false,
+			multiplexed: true,
 		},
 		{
 			name: "multiple frames",
@@ -973,6 +975,7 @@ func TestDecodeDockerLogStream(t *testing.T) {
 			},
 			expected:    "HelloWorld",
 			expectError: false,
+			multiplexed: true,
 		},
 		{
 			name: "zero length frame",
@@ -985,12 +988,20 @@ func TestDecodeDockerLogStream(t *testing.T) {
 			},
 			expected:    "Hello",
 			expectError: false,
+			multiplexed: true,
 		},
 		{
 			name:        "empty input",
 			input:       []byte{},
 			expected:    "",
 			expectError: false,
+			multiplexed: true,
+		},
+		{
+			name:        "raw stream (not multiplexed)",
+			input:       []byte("raw log content"),
+			expected:    "raw log content",
+			multiplexed: false,
 		},
 	}
 
@@ -998,7 +1009,7 @@ func TestDecodeDockerLogStream(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			reader := bytes.NewReader(tt.input)
 			var builder strings.Builder
-			err := decodeDockerLogStream(reader, &builder)
+			err := decodeDockerLogStream(reader, &builder, tt.multiplexed)
 
 			if tt.expectError {
 				assert.Error(t, err)
@@ -1022,7 +1033,7 @@ func TestDecodeDockerLogStreamMemoryProtection(t *testing.T) {
 
 		reader := bytes.NewReader(input)
 		var builder strings.Builder
-		err := decodeDockerLogStream(reader, &builder)
+		err := decodeDockerLogStream(reader, &builder, true)
 
 		assert.Error(t, err)
 		assert.Contains(t, err.Error(), "log frame size")
@@ -1056,7 +1067,7 @@ func TestDecodeDockerLogStreamMemoryProtection(t *testing.T) {
 
 		reader := bytes.NewReader(input)
 		var builder strings.Builder
-		err := decodeDockerLogStream(reader, &builder)
+		err := decodeDockerLogStream(reader, &builder, true)
 
 		// Should complete without error (graceful truncation)
 		assert.NoError(t, err)

agent/gpu.go

@@ -15,7 +15,7 @@ import (
 
 	"github.com/henrygd/beszel/internal/entities/system"
 
-	"golang.org/x/exp/slog"
+	"log/slog"
 )
 
 const (
@@ -44,6 +44,7 @@ type GPUManager struct {
 	rocmSmi       bool
 	tegrastats    bool
 	intelGpuStats bool
+	nvml          bool
 	GpuDataMap    map[string]*system.GPUData
 	// lastAvgData stores the last calculated averages for each GPU
 	// Used when a collection happens before new data arrives (Count == 0)
@@ -297,8 +298,13 @@ func (gm *GPUManager) calculateGPUAverage(id string, gpu *system.GPUData, cacheK
 	currentCount := uint32(gpu.Count)
 	deltaCount := gm.calculateDeltaCount(currentCount, lastSnapshot)
 
-	// If no new data arrived, use last known average
+	// If no new data arrived
 	if deltaCount == 0 {
+		// If GPU appears suspended (instantaneous values are 0), return zero values
+		// Otherwise return last known average for temporary collection gaps
+		if gpu.Temperature == 0 && gpu.MemoryUsed == 0 {
+			return system.GPUData{Name: gpu.Name}
+		}
 		return gm.lastAvgData[id] // zero value if not found
 	}
 
@@ -396,7 +402,7 @@ func (gm *GPUManager) detectGPUs() error {
 	if _, err := exec.LookPath(intelGpuStatsCmd); err == nil {
 		gm.intelGpuStats = true
 	}
-	if gm.nvidiaSmi || gm.rocmSmi || gm.tegrastats || gm.intelGpuStats {
+	if gm.nvidiaSmi || gm.rocmSmi || gm.tegrastats || gm.intelGpuStats || gm.nvml {
 		return nil
 	}
 	return fmt.Errorf("no GPU found - install nvidia-smi, rocm-smi, tegrastats, or intel_gpu_top")
@@ -467,7 +473,20 @@ func NewGPUManager() (*GPUManager, error) {
 	gm.GpuDataMap = make(map[string]*system.GPUData)
 
 	if gm.nvidiaSmi {
-		gm.startCollector(nvidiaSmiCmd)
+		if nvml, _ := GetEnv("NVML"); nvml == "true" {
+			gm.nvml = true
+			gm.nvidiaSmi = false
+			collector := &nvmlCollector{gm: &gm}
+			if err := collector.init(); err == nil {
+				go collector.start()
+			} else {
+				slog.Warn("Failed to initialize NVML, falling back to nvidia-smi", "err", err)
+				gm.nvidiaSmi = true
+				gm.startCollector(nvidiaSmiCmd)
+			}
+		} else {
+			gm.startCollector(nvidiaSmiCmd)
+		}
 	}
 	if gm.rocmSmi {
 		gm.startCollector(rocmSmiCmd)

agent/gpu_nvml.go

@@ -0,0 +1,222 @@
+package agent
+
+import (
+	"fmt"
+	"log/slog"
+	"strings"
+	"time"
+	"unsafe"
+
+	"github.com/ebitengine/purego"
+	"github.com/henrygd/beszel/internal/entities/system"
+)
+
+// NVML constants and types
+const (
+	nvmlSuccess int = 0
+)
+
+type nvmlDevice uintptr
+
+type nvmlReturn int
+
+type nvmlMemoryV1 struct {
+	Total uint64
+	Free  uint64
+	Used  uint64
+}
+
+type nvmlMemoryV2 struct {
+	Version  uint32
+	Total    uint64
+	Reserved uint64
+	Free     uint64
+	Used     uint64
+}
+
+type nvmlUtilization struct {
+	Gpu    uint32
+	Memory uint32
+}
+
+type nvmlPciInfo struct {
+	BusId          [16]byte
+	Domain         uint32
+	Bus            uint32
+	Device         uint32
+	PciDeviceId    uint32
+	PciSubSystemId uint32
+}
+
+// NVML function signatures
+var (
+	nvmlInit                      func() nvmlReturn
+	nvmlShutdown                  func() nvmlReturn
+	nvmlDeviceGetCount            func(count *uint32) nvmlReturn
+	nvmlDeviceGetHandleByIndex    func(index uint32, device *nvmlDevice) nvmlReturn
+	nvmlDeviceGetName             func(device nvmlDevice, name *byte, length uint32) nvmlReturn
+	nvmlDeviceGetMemoryInfo       func(device nvmlDevice, memory uintptr) nvmlReturn
+	nvmlDeviceGetUtilizationRates func(device nvmlDevice, utilization *nvmlUtilization) nvmlReturn
+	nvmlDeviceGetTemperature      func(device nvmlDevice, sensorType int, temp *uint32) nvmlReturn
+	nvmlDeviceGetPowerUsage       func(device nvmlDevice, power *uint32) nvmlReturn
+	nvmlDeviceGetPciInfo          func(device nvmlDevice, pci *nvmlPciInfo) nvmlReturn
+	nvmlErrorString               func(result nvmlReturn) string
+)
+
+type nvmlCollector struct {
+	gm      *GPUManager
+	lib     uintptr
+	devices []nvmlDevice
+	bdfs    []string
+	isV2    bool
+}
+
+func (c *nvmlCollector) init() error {
+	slog.Debug("NVML: Initializing")
+	libPath := getNVMLPath()
+
+	lib, err := openLibrary(libPath)
+	if err != nil {
+		return fmt.Errorf("failed to load %s: %w", libPath, err)
+	}
+	c.lib = lib
+
+	purego.RegisterLibFunc(&nvmlInit, lib, "nvmlInit")
+	purego.RegisterLibFunc(&nvmlShutdown, lib, "nvmlShutdown")
+	purego.RegisterLibFunc(&nvmlDeviceGetCount, lib, "nvmlDeviceGetCount")
+	purego.RegisterLibFunc(&nvmlDeviceGetHandleByIndex, lib, "nvmlDeviceGetHandleByIndex")
+	purego.RegisterLibFunc(&nvmlDeviceGetName, lib, "nvmlDeviceGetName")
+	// Try to get v2 memory info, fallback to v1 if not available
+	if hasSymbol(lib, "nvmlDeviceGetMemoryInfo_v2") {
+		c.isV2 = true
+		purego.RegisterLibFunc(&nvmlDeviceGetMemoryInfo, lib, "nvmlDeviceGetMemoryInfo_v2")
+	} else {
+		purego.RegisterLibFunc(&nvmlDeviceGetMemoryInfo, lib, "nvmlDeviceGetMemoryInfo")
+	}
+	purego.RegisterLibFunc(&nvmlDeviceGetUtilizationRates, lib, "nvmlDeviceGetUtilizationRates")
+	purego.RegisterLibFunc(&nvmlDeviceGetTemperature, lib, "nvmlDeviceGetTemperature")
+	purego.RegisterLibFunc(&nvmlDeviceGetPowerUsage, lib, "nvmlDeviceGetPowerUsage")
+	purego.RegisterLibFunc(&nvmlDeviceGetPciInfo, lib, "nvmlDeviceGetPciInfo")
+	purego.RegisterLibFunc(&nvmlErrorString, lib, "nvmlErrorString")
+
+	if ret := nvmlInit(); ret != nvmlReturn(nvmlSuccess) {
+		return fmt.Errorf("nvmlInit failed: %v", ret)
+	}
+
+	var count uint32
+	if ret := nvmlDeviceGetCount(&count); ret != nvmlReturn(nvmlSuccess) {
+		return fmt.Errorf("nvmlDeviceGetCount failed: %v", ret)
+	}
+
+	for i := uint32(0); i < count; i++ {
+		var device nvmlDevice
+		if ret := nvmlDeviceGetHandleByIndex(i, &device); ret == nvmlReturn(nvmlSuccess) {
+			c.devices = append(c.devices, device)
+			// Get BDF for power state check
+			var pci nvmlPciInfo
+			if ret := nvmlDeviceGetPciInfo(device, &pci); ret == nvmlReturn(nvmlSuccess) {
+				busID := string(pci.BusId[:])
+				if idx := strings.Index(busID, "\x00"); idx != -1 {
+					busID = busID[:idx]
+				}
+				c.bdfs = append(c.bdfs, strings.ToLower(busID))
+			} else {
+				c.bdfs = append(c.bdfs, "")
+			}
+		}
+	}
+
+	return nil
+}
+
+func (c *nvmlCollector) start() {
+	defer nvmlShutdown()
+	ticker := time.Tick(3 * time.Second)
+
+	for range ticker {
+		c.collect()
+	}
+}
+
+func (c *nvmlCollector) collect() {
+	c.gm.Lock()
+	defer c.gm.Unlock()
+
+	for i, device := range c.devices {
+		id := fmt.Sprintf("%d", i)
+		bdf := c.bdfs[i]
+
+		// Update GPUDataMap
+		if _, ok := c.gm.GpuDataMap[id]; !ok {
+			var nameBuf [64]byte
+			if ret := nvmlDeviceGetName(device, &nameBuf[0], 64); ret != nvmlReturn(nvmlSuccess) {
+				continue
+			}
+			name := string(nameBuf[:strings.Index(string(nameBuf[:]), "\x00")])
+			name = strings.TrimPrefix(name, "NVIDIA ")
+			c.gm.GpuDataMap[id] = &system.GPUData{Name: strings.TrimSuffix(name, " Laptop GPU")}
+		}
+		gpu := c.gm.GpuDataMap[id]
+
+		if bdf != "" && !c.isGPUActive(bdf) {
+			slog.Debug("NVML: GPU is suspended, skipping", "bdf", bdf)
+			gpu.Temperature = 0
+			gpu.MemoryUsed = 0
+			continue
+		}
+
+		// Utilization
+		var utilization nvmlUtilization
+		if ret := nvmlDeviceGetUtilizationRates(device, &utilization); ret != nvmlReturn(nvmlSuccess) {
+			slog.Debug("NVML: Utilization failed (GPU likely suspended)", "bdf", bdf, "ret", ret)
+			gpu.Temperature = 0
+			gpu.MemoryUsed = 0
+			continue
+		}
+
+		slog.Debug("NVML: Collecting data for GPU", "bdf", bdf)
+
+		// Temperature
+		var temp uint32
+		nvmlDeviceGetTemperature(device, 0, &temp) // 0 is NVML_TEMPERATURE_GPU
+
+		// Memory: only poll if GPU is active to avoid leaving D3cold state (#1522)
+		if utilization.Gpu > 0 {
+			var usedMem, totalMem uint64
+			if c.isV2 {
+				var memory nvmlMemoryV2
+				memory.Version = 0x02000028 // (2 << 24) | 40 bytes
+				if ret := nvmlDeviceGetMemoryInfo(device, uintptr(unsafe.Pointer(&memory))); ret != nvmlReturn(nvmlSuccess) {
+					slog.Debug("NVML: MemoryInfo_v2 failed", "bdf", bdf, "ret", ret)
+				} else {
+					usedMem = memory.Used
+					totalMem = memory.Total
+				}
+			} else {
+				var memory nvmlMemoryV1
+				if ret := nvmlDeviceGetMemoryInfo(device, uintptr(unsafe.Pointer(&memory))); ret != nvmlReturn(nvmlSuccess) {
+					slog.Debug("NVML: MemoryInfo failed", "bdf", bdf, "ret", ret)
+				} else {
+					usedMem = memory.Used
+					totalMem = memory.Total
+				}
+			}
+			if totalMem > 0 {
+				gpu.MemoryUsed = float64(usedMem) / 1024 / 1024 / mebibytesInAMegabyte
+				gpu.MemoryTotal = float64(totalMem) / 1024 / 1024 / mebibytesInAMegabyte
+			}
+		} else {
+			slog.Debug("NVML: Skipping memory info (utilization=0)", "bdf", bdf)
+		}
+
+		// Power
+		var power uint32
+		nvmlDeviceGetPowerUsage(device, &power)
+
+		gpu.Temperature = float64(temp)
+		gpu.Usage += float64(utilization.Gpu)
+		gpu.Power += float64(power) / 1000.0
+		gpu.Count++
+		slog.Debug("NVML: Collected data", "gpu", gpu)
+	}
+}

agent/gpu_nvml_linux.go

@@ -0,0 +1,57 @@
+//go:build linux
+
+package agent
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/ebitengine/purego"
+	"log/slog"
+)
+
+func openLibrary(name string) (uintptr, error) {
+	return purego.Dlopen(name, purego.RTLD_NOW|purego.RTLD_GLOBAL)
+}
+
+func getNVMLPath() string {
+	return "libnvidia-ml.so.1"
+}
+
+func hasSymbol(lib uintptr, symbol string) bool {
+	_, err := purego.Dlsym(lib, symbol)
+	return err == nil
+}
+
+func (c *nvmlCollector) isGPUActive(bdf string) bool {
+	// runtime_status
+	statusPath := filepath.Join("/sys/bus/pci/devices", bdf, "power/runtime_status")
+	status, err := os.ReadFile(statusPath)
+	if err != nil {
+		slog.Debug("NVML: Can't read runtime_status", "bdf", bdf, "err", err)
+		return true // Assume active if we can't read status
+	}
+	statusStr := strings.TrimSpace(string(status))
+	if statusStr != "active" && statusStr != "resuming" {
+		slog.Debug("NVML: GPU not active", "bdf", bdf, "status", statusStr)
+		return false
+	}
+
+	// power_state (D0 check)
+	// Find any drm card device power_state
+	pstatePathPattern := filepath.Join("/sys/bus/pci/devices", bdf, "drm/card*/device/power_state")
+	matches, _ := filepath.Glob(pstatePathPattern)
+	if len(matches) > 0 {
+		pstate, err := os.ReadFile(matches[0])
+		if err == nil {
+			pstateStr := strings.TrimSpace(string(pstate))
+			if pstateStr != "D0" {
+				slog.Debug("NVML: GPU not in D0 state", "bdf", bdf, "pstate", pstateStr)
+				return false
+			}
+		}
+	}
+
+	return true
+}

agent/gpu_nvml_unsupported.go

@@ -0,0 +1,21 @@
+//go:build !linux && !windows
+
+package agent
+
+import "fmt"
+
+func openLibrary(name string) (uintptr, error) {
+	return 0, fmt.Errorf("nvml not supported on this platform")
+}
+
+func getNVMLPath() string {
+	return ""
+}
+
+func hasSymbol(lib uintptr, symbol string) bool {
+	return false
+}
+
+func (c *nvmlCollector) isGPUActive(bdf string) bool {
+	return true
+}

agent/gpu_nvml_windows.go

@@ -0,0 +1,25 @@
+//go:build windows
+
+package agent
+
+import (
+	"golang.org/x/sys/windows"
+)
+
+func openLibrary(name string) (uintptr, error) {
+	handle, err := windows.LoadLibrary(name)
+	return uintptr(handle), err
+}
+
+func getNVMLPath() string {
+	return "nvml.dll"
+}
+
+func hasSymbol(lib uintptr, symbol string) bool {
+	_, err := windows.GetProcAddress(windows.Handle(lib), symbol)
+	return err == nil
+}
+
+func (c *nvmlCollector) isGPUActive(bdf string) bool {
+	return true
+}

agent/gpu_test.go

@@ -825,7 +825,7 @@ func TestInitializeSnapshots(t *testing.T) {
 }
 
 func TestCalculateGPUAverage(t *testing.T) {
-	t.Run("returns zero value when deltaCount is zero", func(t *testing.T) {
+	t.Run("returns cached average when deltaCount is zero", func(t *testing.T) {
 		gm := &GPUManager{
 			lastSnapshots: map[uint16]map[string]*gpuSnapshot{
 				5000: {
@@ -838,9 +838,10 @@ func TestCalculateGPUAverage(t *testing.T) {
 		}
 
 		gpu := &system.GPUData{
-			Count: 10.0, // Same as snapshot, so delta = 0
-			Usage: 100.0,
-			Power: 200.0,
+			Count:       10.0, // Same as snapshot, so delta = 0
+			Usage:       100.0,
+			Power:       200.0,
+			Temperature: 50.0, // Non-zero to avoid "suspended" check
 		}
 
 		result := gm.calculateGPUAverage("0", gpu, 5000)
@@ -849,6 +850,31 @@ func TestCalculateGPUAverage(t *testing.T) {
 		assert.Equal(t, 100.0, result.Power, "Should return cached average")
 	})
 
+	t.Run("returns zero value when GPU is suspended", func(t *testing.T) {
+		gm := &GPUManager{
+			lastSnapshots: map[uint16]map[string]*gpuSnapshot{
+				5000: {
+					"0": {count: 10, usage: 100, power: 200},
+				},
+			},
+			lastAvgData: map[string]system.GPUData{
+				"0": {Usage: 50.0, Power: 100.0},
+			},
+		}
+
+		gpu := &system.GPUData{
+			Name:        "Test GPU",
+			Count:       10.0,
+			Temperature: 0,
+			MemoryUsed:  0,
+		}
+
+		result := gm.calculateGPUAverage("0", gpu, 5000)
+
+		assert.Equal(t, 0.0, result.Usage, "Should return zero usage")
+		assert.Equal(t, 0.0, result.Power, "Should return zero power")
+	})
+
 	t.Run("calculates average for standard GPU", func(t *testing.T) {
 		gm := &GPUManager{
 			lastSnapshots: map[uint16]map[string]*gpuSnapshot{

agent/handlers.go

@@ -9,7 +9,7 @@ import (
 	"github.com/henrygd/beszel/internal/common"
 	"github.com/henrygd/beszel/internal/entities/smart"
 
-	"golang.org/x/exp/slog"
+	"log/slog"
 )
 
 // HandlerContext provides context for request handlers

agent/response.go

@@ -0,0 +1,31 @@
+package agent
+
+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"
+)
+
+// newAgentResponse creates an AgentResponse using legacy typed fields.
+// This maintains backward compatibility with <= 0.17 hubs that expect specific fields.
+func newAgentResponse(data any, requestID *uint32) common.AgentResponse {
+	response := common.AgentResponse{Id: requestID}
+	switch v := data.(type) {
+	case *system.CombinedData:
+		response.SystemData = v
+	case *common.FingerprintResponse:
+		response.Fingerprint = v
+	case string:
+		response.String = &v
+	case map[string]smart.SmartData:
+		response.SmartData = v
+	case systemd.ServiceDetails:
+		response.ServiceInfo = v
+	default:
+		// For unknown types, use the generic Data field
+		response.Data, _ = cbor.Marshal(data)
+	}
+	return response
+}

agent/server.go

@@ -13,9 +13,7 @@ import (
 
 	"github.com/henrygd/beszel"
 	"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/blang/semver"
 	"github.com/fxamacker/cbor/v2"
@@ -165,20 +163,9 @@ func (a *Agent) handleSSHRequest(w io.Writer, req *common.HubRequest[cbor.RawMes
 	}
 
 	// responder that writes AgentResponse to stdout
+	// Uses legacy typed fields for backward compatibility with <= 0.17
 	sshResponder := func(data any, requestID *uint32) error {
-		response := common.AgentResponse{Id: requestID}
-		switch v := data.(type) {
-		case *system.CombinedData:
-			response.SystemData = v
-		case string:
-			response.String = &v
-		case map[string]smart.SmartData:
-			response.SmartData = v
-		case systemd.ServiceDetails:
-			response.ServiceInfo = v
-		default:
-			response.Error = fmt.Sprintf("unsupported response type: %T", data)
-		}
+		response := newAgentResponse(data, requestID)
 		return cbor.NewEncoder(w).Encode(response)
 	}
 

agent/smart.go

@@ -19,7 +19,7 @@ import (
 
 	"github.com/henrygd/beszel/internal/entities/smart"
 
-	"golang.org/x/exp/slog"
+	"log/slog"
 )
 
 // SmartManager manages data collection for SMART devices

agent/systemd.go

@@ -8,6 +8,7 @@ import (
 	"log/slog"
 	"maps"
 	"math"
+	"os"
 	"strconv"
 	"strings"
 	"sync"
@@ -28,11 +29,29 @@ type systemdManager struct {
 	patterns        []string
 }
 
+// isSystemdAvailable checks if systemd is used on the system to avoid unnecessary connection attempts.
+func isSystemdAvailable() bool {
+	if _, err := os.Stat("/run/systemd/system"); err == nil {
+		return true
+	}
+	if data, err := os.ReadFile("/proc/1/comm"); err == nil {
+		return strings.TrimSpace(string(data)) == "systemd"
+	}
+	return false
+}
+
 // newSystemdManager creates a new systemdManager.
 func newSystemdManager() (*systemdManager, error) {
 	if skipSystemd, _ := GetEnv("SKIP_SYSTEMD"); skipSystemd == "true" {
 		return nil, nil
 	}
+
+	// Check if systemd is available on the system before attempting connection
+	if !isSystemdAvailable() {
+		slog.Debug("Systemd not available on this system")
+		return nil, nil
+	}
+
 	conn, err := dbus.NewSystemConnectionContext(context.Background())
 	if err != nil {
 		slog.Debug("Error connecting to systemd", "err", err, "ref", "https://beszel.dev/guide/systemd")

agent/systemd_test.go

@@ -4,6 +4,7 @@ package agent
 
 import (
 	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -48,6 +49,35 @@ func TestUnescapeServiceNameInvalid(t *testing.T) {
 	}
 }
 
+func TestIsSystemdAvailable(t *testing.T) {
+	// Note: This test's result will vary based on the actual system running the tests
+	// On systems with systemd, it should return true
+	// On systems without systemd, it should return false
+	result := isSystemdAvailable()
+
+	// Check if either the /run/systemd/system directory exists or PID 1 is systemd
+	runSystemdExists := false
+	if _, err := os.Stat("/run/systemd/system"); err == nil {
+		runSystemdExists = true
+	}
+
+	pid1IsSystemd := false
+	if data, err := os.ReadFile("/proc/1/comm"); err == nil {
+		pid1IsSystemd = strings.TrimSpace(string(data)) == "systemd"
+	}
+
+	expected := runSystemdExists || pid1IsSystemd
+
+	assert.Equal(t, expected, result, "isSystemdAvailable should correctly detect systemd presence")
+
+	// Log the result for informational purposes
+	if result {
+		t.Log("Systemd is available on this system")
+	} else {
+		t.Log("Systemd is not available on this system")
+	}
+}
+
 func TestGetServicePatterns(t *testing.T) {
 	tests := []struct {
 		name           string

beszel.go

@@ -6,7 +6,7 @@ import "github.com/blang/semver"
 
 const (
 	// Version is the current version of the application.
-	Version = "0.18.0-beta.1"
+	Version = "0.18.0-beta.2"
 	// AppName is the name of the application.
 	AppName = "beszel"
 )

go.mod

@@ -6,6 +6,7 @@ require (
 	github.com/blang/semver v3.5.1+incompatible
 	github.com/coreos/go-systemd/v22 v22.6.0
 	github.com/distatus/battery v0.11.0
+	github.com/ebitengine/purego v0.9.1
 	github.com/fxamacker/cbor/v2 v2.9.0
 	github.com/gliderlabs/ssh v0.3.8
 	github.com/google/uuid v1.6.0
@@ -20,6 +21,7 @@ require (
 	github.com/stretchr/testify v1.11.1
 	golang.org/x/crypto v0.45.0
 	golang.org/x/exp v0.0.0-20251125195548-87e1e737ad39
+	golang.org/x/sys v0.38.0
 	gopkg.in/yaml.v3 v3.0.1
 )
 
@@ -31,7 +33,6 @@ require (
 	github.com/dolthub/maphash v0.1.0 // indirect
 	github.com/domodwyer/mailyak/v3 v3.6.2 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
-	github.com/ebitengine/purego v0.9.1 // indirect
 	github.com/fatih/color v1.18.0 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.11 // indirect
 	github.com/ganigeorgiev/fexpr v0.5.0 // indirect
@@ -57,7 +58,6 @@ require (
 	golang.org/x/net v0.47.0 // indirect
 	golang.org/x/oauth2 v0.33.0 // indirect
 	golang.org/x/sync v0.18.0 // indirect
-	golang.org/x/sys v0.38.0 // indirect
 	golang.org/x/term v0.37.0 // indirect
 	golang.org/x/text v0.31.0 // indirect
 	howett.net/plist v1.0.1 // indirect

go.sum

@@ -62,8 +62,6 @@ github.com/google/pprof v0.0.0-20251114195745-4902fdda35c8 h1:3DsUAV+VNEQa2CUVLx
 github.com/google/pprof v0.0.0-20251114195745-4902fdda35c8/go.mod h1:I6V7YzU0XDpsHqbsyrghnFZLO1gwK6NPTNvmetQIk9U=
 github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
 github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
-github.com/hashicorp/golang-lru/v2 v2.0.7 h1:a+bsQ5rvGLjzHuww6tVxozPZFVghXaHOwFs4luLUK2k=
-github.com/hashicorp/golang-lru/v2 v2.0.7/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/jarcoal/httpmock v1.4.1 h1:0Ju+VCFuARfFlhVXFc2HxlcQkfB+Xq12/EotHko+x2A=
@@ -173,22 +171,18 @@ gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 howett.net/plist v1.0.1 h1:37GdZ8tP09Q35o9ych3ehygcsL+HqKSwzctveSlarvM=
 howett.net/plist v1.0.1/go.mod h1:lqaXoTrLY4hg8tnEzNru53gicrbv7rrk+2xJA/7hw9g=
-modernc.org/cc/v4 v4.27.1 h1:9W30zRlYrefrDV2JE2O8VDtJ1yPGownxciz5rrbQZis=
-modernc.org/cc/v4 v4.27.1/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
-modernc.org/ccgo/v4 v4.30.1 h1:4r4U1J6Fhj98NKfSjnPUN7Ze2c6MnAdL0hWw6+LrJpc=
-modernc.org/ccgo/v4 v4.30.1/go.mod h1:bIOeI1JL54Utlxn+LwrFyjCx2n2RDiYEaJVSrgdrRfM=
+modernc.org/cc/v4 v4.26.5 h1:xM3bX7Mve6G8K8b+T11ReenJOT+BmVqQj0FY5T4+5Y4=
+modernc.org/cc/v4 v4.26.5/go.mod h1:uVtb5OGqUKpoLWhqwNQo/8LwvoiEBLvZXIQ/SmO6mL0=
+modernc.org/ccgo/v4 v4.28.1 h1:wPKYn5EC/mYTqBO373jKjvX2n+3+aK7+sICCv4Fjy1A=
+modernc.org/ccgo/v4 v4.28.1/go.mod h1:uD+4RnfrVgE6ec9NGguUNdhqzNIeeomeXf6CL0GTE5Q=
 modernc.org/fileutil v1.3.40 h1:ZGMswMNc9JOCrcrakF1HrvmergNLAmxOPjizirpfqBA=
 modernc.org/fileutil v1.3.40/go.mod h1:HxmghZSZVAz/LXcMNwZPA/DRrQZEVP9VX0V4LQGQFOc=
 modernc.org/gc/v2 v2.6.5 h1:nyqdV8q46KvTpZlsw66kWqwXRHdjIlJOhG6kxiV/9xI=
 modernc.org/gc/v2 v2.6.5/go.mod h1:YgIahr1ypgfe7chRuJi2gD7DBQiKSLMPgBQe9oIiito=
-modernc.org/gc/v3 v3.1.1 h1:k8T3gkXWY9sEiytKhcgyiZ2L0DTyCQ/nvX+LoCljoRE=
-modernc.org/gc/v3 v3.1.1/go.mod h1:HFK/6AGESC7Ex+EZJhJ2Gni6cTaYpSMmU/cT9RmlfYY=
 modernc.org/goabi0 v0.2.0 h1:HvEowk7LxcPd0eq6mVOAEMai46V+i7Jrj13t4AzuNks=
 modernc.org/goabi0 v0.2.0/go.mod h1:CEFRnnJhKvWT1c1JTI3Avm+tgOWbkOu5oPA8eH8LnMI=
 modernc.org/libc v1.66.10 h1:yZkb3YeLx4oynyR+iUsXsybsX4Ubx7MQlSYEw4yj59A=
 modernc.org/libc v1.66.10/go.mod h1:8vGSEwvoUoltr4dlywvHqjtAqHBaw0j1jI7iFBTAr2I=
-modernc.org/libc v1.67.1 h1:bFaqOaa5/zbWYJo8aW0tXPX21hXsngG2M7mckCnFSVk=
-modernc.org/libc v1.67.1/go.mod h1:QvvnnJ5P7aitu0ReNpVIEyesuhmDLQ8kaEoyMjIFZJA=
 modernc.org/mathutil v1.7.1 h1:GCZVGXdaN8gTqB1Mf/usp1Y/hSqgI2vAGGP4jZMCxOU=
 modernc.org/mathutil v1.7.1/go.mod h1:4p5IwJITfppl0G4sUEDtCr4DthTaT47/N3aT6MhfgJg=
 modernc.org/memory v1.11.0 h1:o4QC8aMQzmcwCK3t3Ux/ZHmwFPzE6hf2Y5LbkRs+hbI=

internal/common/common-ws.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"github.com/fxamacker/cbor/v2"
 	"github.com/henrygd/beszel/internal/entities/smart"
 	"github.com/henrygd/beszel/internal/entities/system"
 	"github.com/henrygd/beszel/internal/entities/systemd"
@@ -34,14 +35,14 @@ type HubRequest[T any] struct {
 // AgentResponse defines the structure for responses sent from agent to hub.
 type AgentResponse struct {
 	Id          *uint32                    `cbor:"0,keyasint,omitempty"`
-	SystemData  *system.CombinedData       `cbor:"1,keyasint,omitempty,omitzero"`
-	Fingerprint *FingerprintResponse       `cbor:"2,keyasint,omitempty,omitzero"`
+	SystemData  *system.CombinedData       `cbor:"1,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
+	Fingerprint *FingerprintResponse       `cbor:"2,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
 	Error       string                     `cbor:"3,keyasint,omitempty,omitzero"`
-	String      *string                    `cbor:"4,keyasint,omitempty,omitzero"`
-	SmartData   map[string]smart.SmartData `cbor:"5,keyasint,omitempty,omitzero"`
-	ServiceInfo systemd.ServiceDetails     `cbor:"6,keyasint,omitempty,omitzero"`
-	// Logs        *LogsPayload         `cbor:"4,keyasint,omitempty,omitzero"`
-	// RawBytes    []byte               `cbor:"4,keyasint,omitempty,omitzero"`
+	String      *string                    `cbor:"4,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
+	SmartData   map[string]smart.SmartData `cbor:"5,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
+	ServiceInfo systemd.ServiceDetails     `cbor:"6,keyasint,omitempty,omitzero"` // Legacy (<= 0.17)
+	// Data is the generic response payload for new endpoints (0.18+)
+	Data cbor.RawMessage `cbor:"7,keyasint,omitempty,omitzero"`
 }
 
 type FingerprintRequest struct {

internal/hub/systems/system.go

@@ -13,9 +13,11 @@ import (
 	"time"
 
 	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/hub/transport"
 	"github.com/henrygd/beszel/internal/hub/ws"
 
 	"github.com/henrygd/beszel/internal/entities/container"
+	"github.com/henrygd/beszel/internal/entities/smart"
 	"github.com/henrygd/beszel/internal/entities/system"
 	"github.com/henrygd/beszel/internal/entities/systemd"
 
@@ -23,28 +25,30 @@ import (
 
 	"github.com/blang/semver"
 	"github.com/fxamacker/cbor/v2"
+	"github.com/lxzan/gws"
 	"github.com/pocketbase/dbx"
 	"github.com/pocketbase/pocketbase/core"
 	"golang.org/x/crypto/ssh"
 )
 
 type System struct {
-	Id             string               `db:"id"`
-	Host           string               `db:"host"`
-	Port           string               `db:"port"`
-	Status         string               `db:"status"`
-	manager        *SystemManager       // Manager that this system belongs to
-	client         *ssh.Client          // SSH client for fetching data
-	data           *system.CombinedData // system data from agent
-	ctx            context.Context      // Context for stopping the updater
-	cancel         context.CancelFunc   // Stops and removes system from updater
-	WsConn         *ws.WsConn           // Handler for agent WebSocket connection
-	agentVersion   semver.Version       // Agent version
-	updateTicker   *time.Ticker         // Ticker for updating the system
-	detailsFetched atomic.Bool          // True if static system details have been fetched and saved
-	smartFetching  atomic.Bool          // True if SMART devices are currently being fetched
-	smartInterval  time.Duration        // Interval for periodic SMART data updates
-	lastSmartFetch atomic.Int64         // Unix milliseconds of last SMART data fetch
+	Id             string                  `db:"id"`
+	Host           string                  `db:"host"`
+	Port           string                  `db:"port"`
+	Status         string                  `db:"status"`
+	manager        *SystemManager          // Manager that this system belongs to
+	client         *ssh.Client             // SSH client for fetching data
+	sshTransport   *transport.SSHTransport // SSH transport for requests
+	data           *system.CombinedData    // system data from agent
+	ctx            context.Context         // Context for stopping the updater
+	cancel         context.CancelFunc      // Stops and removes system from updater
+	WsConn         *ws.WsConn              // Handler for agent WebSocket connection
+	agentVersion   semver.Version          // Agent version
+	updateTicker   *time.Ticker            // Ticker for updating the system
+	detailsFetched atomic.Bool             // True if static system details have been fetched and saved
+	smartFetching  atomic.Bool             // True if SMART devices are currently being fetched
+	smartInterval  time.Duration           // Interval for periodic SMART data updates
+	lastSmartFetch atomic.Int64            // Unix milliseconds of last SMART data fetch
 }
 
 func (sm *SystemManager) NewSystem(systemId string) *System {
@@ -359,8 +363,78 @@ func (sys *System) getContext() (context.Context, context.CancelFunc) {
 	return sys.ctx, sys.cancel
 }
 
-// fetchDataFromAgent attempts to fetch data from the agent,
-// prioritizing WebSocket if available.
+// request sends a request to the agent, trying WebSocket first, then SSH.
+// This is the unified request method that uses the transport abstraction.
+func (sys *System) request(ctx context.Context, action common.WebSocketAction, req any, dest any) error {
+	// Try WebSocket first
+	if sys.WsConn != nil && sys.WsConn.IsConnected() {
+		wsTransport := transport.NewWebSocketTransport(sys.WsConn)
+		if err := wsTransport.Request(ctx, action, req, dest); err == nil {
+			return nil
+		} else if !shouldFallbackToSSH(err) {
+			return err
+		} else if shouldCloseWebSocket(err) {
+			sys.closeWebSocketConnection()
+		}
+	}
+
+	// Fall back to SSH if WebSocket fails
+	if err := sys.ensureSSHTransport(); err != nil {
+		return err
+	}
+	err := sys.sshTransport.RequestWithRetry(ctx, action, req, dest, 1)
+	// Keep legacy SSH client/version fields in sync for other code paths.
+	if sys.sshTransport != nil {
+		sys.client = sys.sshTransport.GetClient()
+		sys.agentVersion = sys.sshTransport.GetAgentVersion()
+	}
+	return err
+}
+
+func shouldFallbackToSSH(err error) bool {
+	if err == nil {
+		return false
+	}
+	if errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled) {
+		return true
+	}
+	if errors.Is(err, gws.ErrConnClosed) {
+		return true
+	}
+	return errors.Is(err, transport.ErrWebSocketNotConnected)
+}
+
+func shouldCloseWebSocket(err error) bool {
+	if err == nil {
+		return false
+	}
+	return errors.Is(err, gws.ErrConnClosed) || errors.Is(err, transport.ErrWebSocketNotConnected)
+}
+
+// ensureSSHTransport ensures the SSH transport is initialized and connected.
+func (sys *System) ensureSSHTransport() error {
+	if sys.sshTransport == nil {
+		if sys.manager.sshConfig == nil {
+			if err := sys.manager.createSSHClientConfig(); err != nil {
+				return err
+			}
+		}
+		sys.sshTransport = transport.NewSSHTransport(transport.SSHTransportConfig{
+			Host:    sys.Host,
+			Port:    sys.Port,
+			Config:  sys.manager.sshConfig,
+			Timeout: 4 * time.Second,
+		})
+	}
+	// Sync client state with transport
+	if sys.client != nil {
+		sys.sshTransport.SetClient(sys.client)
+		sys.sshTransport.SetAgentVersion(sys.agentVersion)
+	}
+	return nil
+}
+
+// fetchDataFromAgent attempts to fetch data from the agent, prioritizing WebSocket if available.
 func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*system.CombinedData, error) {
 	if sys.data == nil {
 		sys.data = &system.CombinedData{}
@@ -386,112 +460,47 @@ func (sys *System) fetchDataViaWebSocket(options common.DataRequestOptions) (*sy
 	if sys.WsConn == nil || !sys.WsConn.IsConnected() {
 		return nil, errors.New("no websocket connection")
 	}
-	err := sys.WsConn.RequestSystemData(context.Background(), sys.data, options)
+	wsTransport := transport.NewWebSocketTransport(sys.WsConn)
+	err := wsTransport.Request(context.Background(), common.GetData, options, sys.data)
 	if err != nil {
 		return nil, err
 	}
 	return sys.data, nil
 }
 
-// fetchStringFromAgentViaSSH is a generic function to fetch strings via SSH
-func (sys *System) fetchStringFromAgentViaSSH(action common.WebSocketAction, requestData any, errorMsg string) (string, error) {
-	var result string
-	err := sys.runSSHOperation(4*time.Second, 1, func(session *ssh.Session) (bool, error) {
-		stdout, err := session.StdoutPipe()
-		if err != nil {
-			return false, err
-		}
-		stdin, stdinErr := session.StdinPipe()
-		if stdinErr != nil {
-			return false, stdinErr
-		}
-		if err := session.Shell(); err != nil {
-			return false, err
-		}
-		req := common.HubRequest[any]{Action: action, Data: requestData}
-		_ = cbor.NewEncoder(stdin).Encode(req)
-		_ = stdin.Close()
-		var resp common.AgentResponse
-		err = cbor.NewDecoder(stdout).Decode(&resp)
-		if err != nil {
-			return false, err
-		}
-		if resp.String == nil {
-			return false, errors.New(errorMsg)
-		}
-		result = *resp.String
-		return false, nil
-	})
-	return result, err
-}
-
 // FetchContainerInfoFromAgent fetches container info from the agent
 func (sys *System) FetchContainerInfoFromAgent(containerID string) (string, error) {
-	// fetch via websocket
-	if sys.WsConn != nil && sys.WsConn.IsConnected() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		return sys.WsConn.RequestContainerInfo(ctx, containerID)
-	}
-	// fetch via SSH
-	return sys.fetchStringFromAgentViaSSH(common.GetContainerInfo, common.ContainerInfoRequest{ContainerID: containerID}, "no info in response")
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	var result string
+	err := sys.request(ctx, common.GetContainerInfo, common.ContainerInfoRequest{ContainerID: containerID}, &result)
+	return result, err
 }
 
 // FetchContainerLogsFromAgent fetches container logs from the agent
 func (sys *System) FetchContainerLogsFromAgent(containerID string) (string, error) {
-	// fetch via websocket
-	if sys.WsConn != nil && sys.WsConn.IsConnected() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		return sys.WsConn.RequestContainerLogs(ctx, containerID)
-	}
-	// fetch via SSH
-	return sys.fetchStringFromAgentViaSSH(common.GetContainerLogs, common.ContainerLogsRequest{ContainerID: containerID}, "no logs in response")
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	var result string
+	err := sys.request(ctx, common.GetContainerLogs, common.ContainerLogsRequest{ContainerID: containerID}, &result)
+	return result, err
 }
 
 // FetchSystemdInfoFromAgent fetches detailed systemd service information from the agent
 func (sys *System) FetchSystemdInfoFromAgent(serviceName string) (systemd.ServiceDetails, error) {
-	// fetch via websocket
-	if sys.WsConn != nil && sys.WsConn.IsConnected() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		return sys.WsConn.RequestSystemdInfo(ctx, serviceName)
-	}
-
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
 	var result systemd.ServiceDetails
-	err := sys.runSSHOperation(5*time.Second, 1, func(session *ssh.Session) (bool, error) {
-		stdout, err := session.StdoutPipe()
-		if err != nil {
-			return false, err
-		}
-		stdin, stdinErr := session.StdinPipe()
-		if stdinErr != nil {
-			return false, stdinErr
-		}
-		if err := session.Shell(); err != nil {
-			return false, err
-		}
-
-		req := common.HubRequest[any]{Action: common.GetSystemdInfo, Data: common.SystemdInfoRequest{ServiceName: serviceName}}
-		if err := cbor.NewEncoder(stdin).Encode(req); err != nil {
-			return false, err
-		}
-		_ = stdin.Close()
-
-		var resp common.AgentResponse
-		if err := cbor.NewDecoder(stdout).Decode(&resp); err != nil {
-			return false, err
-		}
-		if resp.ServiceInfo == nil {
-			if resp.Error != "" {
-				return false, errors.New(resp.Error)
-			}
-			return false, errors.New("no systemd info in response")
-		}
-		result = resp.ServiceInfo
-		return false, nil
-	})
+	err := sys.request(ctx, common.GetSystemdInfo, common.SystemdInfoRequest{ServiceName: serviceName}, &result)
+	return result, err
+}
 
+// FetchSmartDataFromAgent fetches SMART data from the agent
+func (sys *System) FetchSmartDataFromAgent() (map[string]smart.SmartData, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+	var result map[string]smart.SmartData
+	err := sys.request(ctx, common.GetSmartData, nil, &result)
 	return result, err
 }
 
@@ -656,6 +665,9 @@ func (sys *System) createSessionWithTimeout(timeout time.Duration) (*ssh.Session
 
 // closeSSHConnection closes the SSH connection but keeps the system in the manager
 func (sys *System) closeSSHConnection() {
+	if sys.sshTransport != nil {
+		sys.sshTransport.Close()
+	}
 	if sys.client != nil {
 		sys.client.Close()
 		sys.client = nil

internal/hub/systems/system_smart.go

@@ -1,54 +1,14 @@
 package systems
 
 import (
-	"context"
 	"database/sql"
 	"errors"
 	"strings"
-	"time"
 
-	"github.com/fxamacker/cbor/v2"
-	"github.com/henrygd/beszel/internal/common"
 	"github.com/henrygd/beszel/internal/entities/smart"
 	"github.com/pocketbase/pocketbase/core"
-	"golang.org/x/crypto/ssh"
 )
 
-// FetchSmartDataFromAgent fetches SMART data from the agent
-func (sys *System) FetchSmartDataFromAgent() (map[string]smart.SmartData, error) {
-	// fetch via websocket
-	if sys.WsConn != nil && sys.WsConn.IsConnected() {
-		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-		defer cancel()
-		return sys.WsConn.RequestSmartData(ctx)
-	}
-	// fetch via SSH
-	var result map[string]smart.SmartData
-	err := sys.runSSHOperation(5*time.Second, 1, func(session *ssh.Session) (bool, error) {
-		stdout, err := session.StdoutPipe()
-		if err != nil {
-			return false, err
-		}
-		stdin, stdinErr := session.StdinPipe()
-		if stdinErr != nil {
-			return false, stdinErr
-		}
-		if err := session.Shell(); err != nil {
-			return false, err
-		}
-		req := common.HubRequest[any]{Action: common.GetSmartData}
-		_ = cbor.NewEncoder(stdin).Encode(req)
-		_ = stdin.Close()
-		var resp common.AgentResponse
-		if err := cbor.NewDecoder(stdout).Decode(&resp); err != nil {
-			return false, err
-		}
-		result = resp.SmartData
-		return false, nil
-	})
-	return result, err
-}
-
 // FetchAndSaveSmartDevices fetches SMART data from the agent and saves it to the database
 func (sys *System) FetchAndSaveSmartDevices() error {
 	smartData, err := sys.FetchSmartDataFromAgent()

internal/hub/transport/ssh.go

@@ -0,0 +1,227 @@
+package transport
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net"
+	"strings"
+	"time"
+
+	"github.com/blang/semver"
+	"github.com/fxamacker/cbor/v2"
+	"github.com/henrygd/beszel/internal/common"
+	"golang.org/x/crypto/ssh"
+)
+
+// SSHTransport implements Transport over SSH connections.
+type SSHTransport struct {
+	client       *ssh.Client
+	config       *ssh.ClientConfig
+	host         string
+	port         string
+	agentVersion semver.Version
+	timeout      time.Duration
+}
+
+// SSHTransportConfig holds configuration for creating an SSH transport.
+type SSHTransportConfig struct {
+	Host         string
+	Port         string
+	Config       *ssh.ClientConfig
+	AgentVersion semver.Version
+	Timeout      time.Duration
+}
+
+// NewSSHTransport creates a new SSH transport with the given configuration.
+func NewSSHTransport(cfg SSHTransportConfig) *SSHTransport {
+	timeout := cfg.Timeout
+	if timeout == 0 {
+		timeout = 4 * time.Second
+	}
+	return &SSHTransport{
+		config:       cfg.Config,
+		host:         cfg.Host,
+		port:         cfg.Port,
+		agentVersion: cfg.AgentVersion,
+		timeout:      timeout,
+	}
+}
+
+// SetClient sets the SSH client for reuse across requests.
+func (t *SSHTransport) SetClient(client *ssh.Client) {
+	t.client = client
+}
+
+// SetAgentVersion sets the agent version (extracted from SSH handshake).
+func (t *SSHTransport) SetAgentVersion(version semver.Version) {
+	t.agentVersion = version
+}
+
+// GetClient returns the current SSH client (for connection management).
+func (t *SSHTransport) GetClient() *ssh.Client {
+	return t.client
+}
+
+// GetAgentVersion returns the agent version.
+func (t *SSHTransport) GetAgentVersion() semver.Version {
+	return t.agentVersion
+}
+
+// Request sends a request to the agent via SSH and unmarshals the response.
+func (t *SSHTransport) Request(ctx context.Context, action common.WebSocketAction, req any, dest any) error {
+	if t.client == nil {
+		if err := t.connect(); err != nil {
+			return err
+		}
+	}
+
+	session, err := t.createSessionWithTimeout(ctx)
+	if err != nil {
+		return err
+	}
+	defer session.Close()
+
+	stdout, err := session.StdoutPipe()
+	if err != nil {
+		return err
+	}
+	stdin, err := session.StdinPipe()
+	if err != nil {
+		return err
+	}
+	if err := session.Shell(); err != nil {
+		return err
+	}
+
+	// Send request
+	hubReq := common.HubRequest[any]{Action: action, Data: req}
+	if err := cbor.NewEncoder(stdin).Encode(hubReq); err != nil {
+		return fmt.Errorf("failed to encode request: %w", err)
+	}
+	stdin.Close()
+
+	// Read response
+	var resp common.AgentResponse
+	if err := cbor.NewDecoder(stdout).Decode(&resp); err != nil {
+		return fmt.Errorf("failed to decode response: %w", err)
+	}
+
+	if resp.Error != "" {
+		return errors.New(resp.Error)
+	}
+
+	if err := session.Wait(); err != nil {
+		return err
+	}
+
+	return UnmarshalResponse(resp, action, dest)
+}
+
+// IsConnected returns true if the SSH connection is active.
+func (t *SSHTransport) IsConnected() bool {
+	return t.client != nil
+}
+
+// Close terminates the SSH connection.
+func (t *SSHTransport) Close() {
+	if t.client != nil {
+		t.client.Close()
+		t.client = nil
+	}
+}
+
+// connect establishes a new SSH connection.
+func (t *SSHTransport) connect() error {
+	if t.config == nil {
+		return errors.New("SSH config not set")
+	}
+
+	network := "tcp"
+	host := t.host
+	if strings.HasPrefix(host, "/") {
+		network = "unix"
+	} else {
+		host = net.JoinHostPort(host, t.port)
+	}
+
+	client, err := ssh.Dial(network, host, t.config)
+	if err != nil {
+		return err
+	}
+	t.client = client
+
+	// Extract agent version from server version string
+	t.agentVersion, _ = extractAgentVersion(string(client.Conn.ServerVersion()))
+	return nil
+}
+
+// createSessionWithTimeout creates a new SSH session with a timeout.
+func (t *SSHTransport) createSessionWithTimeout(ctx context.Context) (*ssh.Session, error) {
+	if t.client == nil {
+		return nil, errors.New("client not initialized")
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, t.timeout)
+	defer cancel()
+
+	sessionChan := make(chan *ssh.Session, 1)
+	errChan := make(chan error, 1)
+
+	go func() {
+		session, err := t.client.NewSession()
+		if err != nil {
+			errChan <- err
+		} else {
+			sessionChan <- session
+		}
+	}()
+
+	select {
+	case session := <-sessionChan:
+		return session, nil
+	case err := <-errChan:
+		return nil, err
+	case <-ctx.Done():
+		return nil, errors.New("timeout creating session")
+	}
+}
+
+// extractAgentVersion extracts the beszel version from SSH server version string.
+func extractAgentVersion(versionString string) (semver.Version, error) {
+	_, after, _ := strings.Cut(versionString, "_")
+	return semver.Parse(after)
+}
+
+// RequestWithRetry sends a request with automatic retry on connection failures.
+func (t *SSHTransport) RequestWithRetry(ctx context.Context, action common.WebSocketAction, req any, dest any, retries int) error {
+	var lastErr error
+	for attempt := 0; attempt <= retries; attempt++ {
+		err := t.Request(ctx, action, req, dest)
+		if err == nil {
+			return nil
+		}
+		lastErr = err
+
+		// Check if it's a connection error that warrants a retry
+		if isConnectionError(err) && attempt < retries {
+			t.Close()
+			continue
+		}
+		return err
+	}
+	return lastErr
+}
+
+// isConnectionError checks if an error indicates a connection problem.
+func isConnectionError(err error) bool {
+	if err == nil {
+		return false
+	}
+	errStr := err.Error()
+	return strings.Contains(errStr, "connection") ||
+		strings.Contains(errStr, "EOF") ||
+		strings.Contains(errStr, "closed") ||
+		errors.Is(err, io.EOF)
+}

internal/hub/transport/transport.go

@@ -0,0 +1,112 @@
+// Package transport provides a unified abstraction for hub-agent communication
+// over different transports (WebSocket, SSH).
+package transport
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"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"
+)
+
+// Transport defines the interface for hub-agent communication.
+// Both WebSocket and SSH transports implement this interface.
+type Transport interface {
+	// Request sends a request to the agent and unmarshals the response into dest.
+	// The dest parameter should be a pointer to the expected response type.
+	Request(ctx context.Context, action common.WebSocketAction, req any, dest any) error
+	// IsConnected returns true if the transport connection is active.
+	IsConnected() bool
+	// Close terminates the transport connection.
+	Close()
+}
+
+// UnmarshalResponse unmarshals an AgentResponse into the destination type.
+// It first checks the generic Data field (0.19+ agents), then falls back
+// to legacy typed fields for backward compatibility with 0.18.0 agents.
+func UnmarshalResponse(resp common.AgentResponse, action common.WebSocketAction, dest any) error {
+	if dest == nil {
+		return errors.New("nil destination")
+	}
+	// Try generic Data field first (0.19+)
+	if len(resp.Data) > 0 {
+		if err := cbor.Unmarshal(resp.Data, dest); err != nil {
+			return fmt.Errorf("failed to unmarshal generic response data: %w", err)
+		}
+		return nil
+	}
+	// Fall back to legacy typed fields for older agents/hubs.
+	return unmarshalLegacyResponse(resp, action, dest)
+}
+
+// unmarshalLegacyResponse handles legacy responses that use typed fields.
+func unmarshalLegacyResponse(resp common.AgentResponse, action common.WebSocketAction, dest any) error {
+	switch action {
+	case common.GetData:
+		d, ok := dest.(*system.CombinedData)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for GetData: %T", dest)
+		}
+		if resp.SystemData == nil {
+			return errors.New("no system data in response")
+		}
+		*d = *resp.SystemData
+		return nil
+	case common.CheckFingerprint:
+		d, ok := dest.(*common.FingerprintResponse)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for CheckFingerprint: %T", dest)
+		}
+		if resp.Fingerprint == nil {
+			return errors.New("no fingerprint in response")
+		}
+		*d = *resp.Fingerprint
+		return nil
+	case common.GetContainerLogs:
+		d, ok := dest.(*string)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for GetContainerLogs: %T", dest)
+		}
+		if resp.String == nil {
+			return errors.New("no logs in response")
+		}
+		*d = *resp.String
+		return nil
+	case common.GetContainerInfo:
+		d, ok := dest.(*string)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for GetContainerInfo: %T", dest)
+		}
+		if resp.String == nil {
+			return errors.New("no info in response")
+		}
+		*d = *resp.String
+		return nil
+	case common.GetSmartData:
+		d, ok := dest.(*map[string]smart.SmartData)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for GetSmartData: %T", dest)
+		}
+		if resp.SmartData == nil {
+			return errors.New("no SMART data in response")
+		}
+		*d = resp.SmartData
+		return nil
+	case common.GetSystemdInfo:
+		d, ok := dest.(*systemd.ServiceDetails)
+		if !ok {
+			return fmt.Errorf("unexpected dest type for GetSystemdInfo: %T", dest)
+		}
+		if resp.ServiceInfo == nil {
+			return errors.New("no systemd info in response")
+		}
+		*d = resp.ServiceInfo
+		return nil
+	}
+	return fmt.Errorf("unsupported action: %d", action)
+}

internal/hub/transport/websocket.go

@@ -0,0 +1,74 @@
+package transport
+
+import (
+	"context"
+	"errors"
+
+	"github.com/fxamacker/cbor/v2"
+	"github.com/henrygd/beszel"
+	"github.com/henrygd/beszel/internal/common"
+	"github.com/henrygd/beszel/internal/hub/ws"
+)
+
+// ErrWebSocketNotConnected indicates a WebSocket transport is not currently connected.
+var ErrWebSocketNotConnected = errors.New("websocket not connected")
+
+// WebSocketTransport implements Transport over WebSocket connections.
+type WebSocketTransport struct {
+	wsConn *ws.WsConn
+}
+
+// NewWebSocketTransport creates a new WebSocket transport wrapper.
+func NewWebSocketTransport(wsConn *ws.WsConn) *WebSocketTransport {
+	return &WebSocketTransport{wsConn: wsConn}
+}
+
+// Request sends a request to the agent via WebSocket and unmarshals the response.
+func (t *WebSocketTransport) Request(ctx context.Context, action common.WebSocketAction, req any, dest any) error {
+	if !t.IsConnected() {
+		return ErrWebSocketNotConnected
+	}
+
+	pendingReq, err := t.wsConn.SendRequest(ctx, action, req)
+	if err != nil {
+		return err
+	}
+
+	// Wait for response
+	select {
+	case message := <-pendingReq.ResponseCh:
+		defer message.Close()
+		defer pendingReq.Cancel()
+
+		// Legacy agents (< MinVersionAgentResponse) respond with a raw payload instead of an AgentResponse wrapper.
+		if t.wsConn.AgentVersion().LT(beszel.MinVersionAgentResponse) {
+			return cbor.Unmarshal(message.Data.Bytes(), dest)
+		}
+
+		var agentResponse common.AgentResponse
+		if err := cbor.Unmarshal(message.Data.Bytes(), &agentResponse); err != nil {
+			return err
+		}
+
+		if agentResponse.Error != "" {
+			return errors.New(agentResponse.Error)
+		}
+
+		return UnmarshalResponse(agentResponse, action, dest)
+
+	case <-pendingReq.Context.Done():
+		return pendingReq.Context.Err()
+	}
+}
+
+// IsConnected returns true if the WebSocket connection is active.
+func (t *WebSocketTransport) IsConnected() bool {
+	return t.wsConn != nil && t.wsConn.IsConnected()
+}
+
+// Close terminates the WebSocket connection.
+func (t *WebSocketTransport) Close() {
+	if t.wsConn != nil {
+		t.wsConn.Close(nil)
+	}
+}

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

internal/site/package.json

@@ -1,7 +1,7 @@
 {
 	"name": "beszel",
 	"private": true,
-	"version": "0.18.0-beta.1",
+	"version": "0.18.0-beta.2",
 	"type": "module",
 	"scripts": {
 		"dev": "vite --host",

internal/site/src/components/routes/system.tsx

@@ -16,7 +16,7 @@ import MemChart from "@/components/charts/mem-chart"
 import SwapChart from "@/components/charts/swap-chart"
 import TemperatureChart from "@/components/charts/temperature-chart"
 import { getPbTimestamp, pb } from "@/lib/api"
-import { ChartType, Os, SystemStatus, Unit } from "@/lib/enums"
+import { ChartType, SystemStatus, Unit } from "@/lib/enums"
 import { batteryStateTranslations } from "@/lib/i18n"
 import {
 	$allSystemsById,
@@ -222,7 +222,6 @@ export default memo(function SystemDetail({ id }: { id: string }) {
 	}, [system.id])
 
 	// subscribe to realtime metrics if chart time is 1m
-	// biome-ignore lint/correctness/useExhaustiveDependencies: not necessary
 	useEffect(() => {
 		let unsub = () => {}
 		if (!system.id || chartTime !== "1m") {
@@ -260,7 +259,6 @@ export default memo(function SystemDetail({ id }: { id: string }) {
 		}
 	}, [chartTime, system.id])
 
-	// biome-ignore lint/correctness/useExhaustiveDependencies: not necessary
 	const chartData: ChartData = useMemo(() => {
 		const lastCreated = Math.max(
 			(systemStats.at(-1)?.created as number) ?? 0,
@@ -300,7 +298,6 @@ export default memo(function SystemDetail({ id }: { id: string }) {
 	}, [])
 
 	// get stats
-	// biome-ignore lint/correctness/useExhaustiveDependencies: not necessary
 	useEffect(() => {
 		if (!system.id || !chartTime || chartTime === "1m") {
 			return
@@ -403,11 +400,35 @@ export default memo(function SystemDetail({ id }: { id: string }) {
 	const containerFilterBar = containerData.length ? <FilterBar /> : null
 
 	const dataEmpty = !chartLoading && chartData.systemStats.length === 0
-	const lastGpuVals = Object.values(systemStats.at(-1)?.stats.g ?? {})
-	const hasGpuData = lastGpuVals.length > 0
-	const hasGpuPowerData = lastGpuVals.some((gpu) => gpu.p !== undefined || gpu.pp !== undefined)
-	const hasGpuEnginesData = lastGpuVals.some((gpu) => gpu.e !== undefined)
-	const isLinux = (details?.os ?? system.info?.os) === Os.Linux
+	const lastGpus = systemStats.at(-1)?.stats?.g
+
+	let hasGpuData = false
+	let hasGpuEnginesData = false
+	let hasGpuPowerData = false
+
+	if (lastGpus) {
+		// check if there are any GPUs with engines
+		for (const id in lastGpus) {
+			hasGpuData = true
+			if (lastGpus[id].e !== undefined) {
+				hasGpuEnginesData = true
+				break
+			}
+		}
+		// check if there are any GPUs with power data
+		for (let i = 0; i < systemStats.length && !hasGpuPowerData; i++) {
+			const gpus = systemStats[i].stats?.g
+			if (!gpus) continue
+			for (const id in gpus) {
+				if (gpus[id].p !== undefined || gpus[id].pp !== undefined) {
+					hasGpuPowerData = true
+					break
+				}
+			}
+		}
+	}
+
+	const isLinux = !(details?.os ?? system.info?.os)
 	const isPodman = details?.podman ?? system.info?.p ?? false
 
 	return (
@@ -718,64 +739,65 @@ export default memo(function SystemDetail({ id }: { id: string }) {
 								<GpuEnginesChart chartData={chartData} />
 							</ChartCard>
 						)}
-						{Object.keys(systemStats.at(-1)?.stats.g ?? {}).map((id) => {
-							const gpu = systemStats.at(-1)?.stats.g?.[id] as GPUData
-							return (
-								<div key={id} className="contents">
-									<ChartCard
-										className={cn(grid && "!col-span-1")}
-										empty={dataEmpty}
-										grid={grid}
-										title={`${gpu.n} ${t`Usage`}`}
-										description={t`Average utilization of ${gpu.n}`}
-									>
-										<AreaChartDefault
-											chartData={chartData}
-											dataPoints={[
-												{
-													label: t`Usage`,
-													dataKey: ({ stats }) => stats?.g?.[id]?.u ?? 0,
-													color: 1,
-													opacity: 0.35,
-												},
-											]}
-											tickFormatter={(val) => `${toFixedFloat(val, 2)}%`}
-											contentFormatter={({ value }) => `${decimalString(value)}%`}
-										/>
-									</ChartCard>
-
-									{(gpu.mt ?? 0) > 0 && (
+						{lastGpus &&
+							Object.keys(lastGpus).map((id) => {
+								const gpu = lastGpus[id] as GPUData
+								return (
+									<div key={id} className="contents">
 										<ChartCard
+											className={cn(grid && "!col-span-1")}
 											empty={dataEmpty}
 											grid={grid}
-											title={`${gpu.n} VRAM`}
-											description={t`Precise utilization at the recorded time`}
+											title={`${gpu.n} ${t`Usage`}`}
+											description={t`Average utilization of ${gpu.n}`}
 										>
 											<AreaChartDefault
 												chartData={chartData}
 												dataPoints={[
 													{
 														label: t`Usage`,
-														dataKey: ({ stats }) => stats?.g?.[id]?.mu ?? 0,
-														color: 2,
-														opacity: 0.25,
+														dataKey: ({ stats }) => stats?.g?.[id]?.u ?? 0,
+														color: 1,
+														opacity: 0.35,
 													},
 												]}
-												max={gpu.mt}
-												tickFormatter={(val) => {
-													const { value, unit } = formatBytes(val, false, Unit.Bytes, true)
-													return `${toFixedFloat(value, value >= 10 ? 0 : 1)} ${unit}`
-												}}
-												contentFormatter={({ value }) => {
-													const { value: convertedValue, unit } = formatBytes(value, false, Unit.Bytes, true)
-													return `${decimalString(convertedValue)} ${unit}`
-												}}
+												tickFormatter={(val) => `${toFixedFloat(val, 2)}%`}
+												contentFormatter={({ value }) => `${decimalString(value)}%`}
 											/>
 										</ChartCard>
-									)}
-								</div>
-							)
-						})}
+
+										{(gpu.mt ?? 0) > 0 && (
+											<ChartCard
+												empty={dataEmpty}
+												grid={grid}
+												title={`${gpu.n} VRAM`}
+												description={t`Precise utilization at the recorded time`}
+											>
+												<AreaChartDefault
+													chartData={chartData}
+													dataPoints={[
+														{
+															label: t`Usage`,
+															dataKey: ({ stats }) => stats?.g?.[id]?.mu ?? 0,
+															color: 2,
+															opacity: 0.25,
+														},
+													]}
+													max={gpu.mt}
+													tickFormatter={(val) => {
+														const { value, unit } = formatBytes(val, false, Unit.Bytes, true)
+														return `${toFixedFloat(value, value >= 10 ? 0 : 1)} ${unit}`
+													}}
+													contentFormatter={({ value }) => {
+														const { value: convertedValue, unit } = formatBytes(value, false, Unit.Bytes, true)
+														return `${decimalString(convertedValue)} ${unit}`
+													}}
+												/>
+											</ChartCard>
+										)}
+									</div>
+								)
+							})}
 					</div>
 				)}
 

internal/site/src/components/systems-table/systems-table-columns.tsx

@@ -287,12 +287,12 @@ export function SystemsTableColumns(viewMode: "table" | "grid"): ColumnDef<Syste
 					return null
 				}
 
-				const iconColor = pct < 10 ? "text-red-500" : pct < 25 ? "text-yellow-500" : "text-muted-foreground"
-
 				let Icon = PlugChargingIcon
+				let iconColor = "text-muted-foreground"
 
 				if (state !== BatteryState.Charging) {
 					if (pct < 25) {
+						iconColor = pct < 11 ? "text-red-500" : "text-yellow-500"
 						Icon = BatteryLowIcon
 					} else if (pct < 75) {
 						Icon = BatteryMediumIcon

internal/site/src/components/ui/icons.tsx

@@ -154,7 +154,7 @@ export function BatteryMediumIcon(props: SVGProps<SVGSVGElement>) {
 export function BatteryLowIcon(props: SVGProps<SVGSVGElement>) {
 	return (
 		<svg viewBox="0 0 24 24" {...props} fill="currentColor">
-			<path d="M16 20H8V6h8m.67-2H15V2H9v2H7.33C6.6 4 6 4.6 6 5.33v15.34C6 21.4 6.6 22 7.33 22h9.34c.74 0 1.33-.59 1.33-1.33V5.33C18 4.6 17.4 4 16.67 4M15 16H9v3h6zm0-4.5H9v3h6z" />
+			<path d="M16 17H8V6h8m.7-2H15V2H9v2H7.3A1.3 1.3 0 0 0 6 5.3v15.4q.1 1.2 1.3 1.3h9.4a1.3 1.3 0 0 0 1.3-1.3V5.3q-.1-1.2-1.3-1.3" />
 		</svg>
 	)
 }