fix: handle missing results, topic validation, and fragmented WebSocket frames

- LLDP topology and SNMP jobs now always send a result to the server,
  preventing indefinite waits when errors or cancellation occur
- handleMessage validates message topic to reject cross-agent messages
- WebSocket client reassembles fragmented (continuation) frames per RFC 6455
- Replace fragile sync.Once reset in sslRootCAs with mutex-guarded nil check
- Guard nextBackoff against zero/negative input that would panic
- Drain HTTP response body on regex compilation failure for connection reuse

agent.go

@@ -334,7 +334,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
 				slog.Debug("invalid message", "error", err)
 				continue
 			}
-			shouldEnd, endErr := handleMessage(sessionCtx, msg, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh)
+			shouldEnd, endErr := handleMessage(sessionCtx, msg, topic, pools, snmpResultCh, mikrotikResultCh, credTestResultCh, monitoringCheckCh, checkResultCh, lldpTopologyResultCh)
 			if shouldEnd {
 				flushSnmpBatch()
 				return endErr
@@ -401,6 +401,7 @@ func runSession(ctx context.Context, baseURL, token string) error {
 func handleMessage(
 	ctx context.Context,
 	msg channelMsg,
+	topic string,
 	pools *jobPools,
 	snmpResultCh chan<- *pb.SnmpResult,
 	mikrotikResultCh chan<- *pb.MikrotikResult,
@@ -409,6 +410,12 @@ func handleMessage(
 	checkResultCh chan<- *pb.CheckResult,
 	lldpTopologyResultCh chan<- *pb.LldpTopologyResult,
 ) (bool, error) {
+	// Ignore messages not addressed to our topic (except Phoenix control messages)
+	if msg.Topic != topic && msg.Topic != "phoenix" {
+		slog.Debug("ignoring message for different topic", "got", msg.Topic, "want", topic)
+		return false, nil
+	}
+
 	switch msg.Event {
 	case "phx_reply":
 		slog.Debug("channel reply", "topic", msg.Topic)
@@ -539,11 +546,18 @@ func dispatchJob(
 
 // nextBackoff doubles the current delay (capped at max) and adds up to 25% jitter.
 func nextBackoff(current, maxDelay time.Duration) time.Duration {
+	if current <= 0 {
+		current = initialRetryDelay
+	}
 	next := current * 2
 	if next > maxDelay {
 		next = maxDelay
 	}
-	jitter := time.Duration(rand.Int64N(int64(next / 4)))
+	jitterRange := int64(next / 4)
+	if jitterRange <= 0 {
+		return next
+	}
+	jitter := time.Duration(rand.Int64N(jitterRange))
 	return next + jitter
 }
 

agent_test.go

@@ -96,7 +96,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "phx_reply", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "phx_reply", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Just verify it doesn't panic
 	})
 
@@ -124,7 +124,7 @@ func TestHandleMessage(t *testing.T) {
 			SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
 		})
 
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Wait for goroutine to finish
 		select {
 		case <-snmpCh:
@@ -139,7 +139,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: json.RawMessage(`not json`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should log error but not panic
 	})
 
@@ -150,7 +150,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 	})
 
 	t.Run("invalid protobuf", func(t *testing.T) {
@@ -160,7 +160,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 	})
 
 	t.Run("restart", func(t *testing.T) {
@@ -169,7 +169,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		shouldEnd, reason := handleMessage(context.Background(), channelMsg{Event: "restart", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		shouldEnd, reason := handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "restart", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 
 		if !shouldEnd {
 			t.Error("expected handleMessage to return true for restart")
@@ -192,6 +192,7 @@ func TestHandleMessage(t *testing.T) {
 			shouldEnd, reason := handleMessage(
 				context.Background(),
 				channelMsg{Topic: "agent:test", Event: event, Payload: json.RawMessage(`{}`)},
+				"agent:test",
 				testPools(t),
 				snmpCh,
 				mtCh,
@@ -226,7 +227,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 
 		if calledURL != "https://example.com/agent" {
 			t.Errorf("expected update URL %q, got %q", "https://example.com/agent", calledURL)
@@ -250,7 +251,7 @@ func TestHandleMessage(t *testing.T) {
 		checkCh := make(chan *pb.CheckResult, 1)
 		// Missing URL field
 		payload, _ := json.Marshal(map[string]string{"checksum": "abc123"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 
 		if called {
 			t.Error("selfUpdate should not be called with empty URL")
@@ -271,7 +272,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent", "checksum": "abc123"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should log error but not panic
 	})
 
@@ -291,7 +292,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"url": "https://example.com/agent"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "update", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "update", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 
 		if called {
 			t.Error("selfUpdate should not be called with empty checksum")
@@ -304,7 +305,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "some_unknown_event", Payload: json.RawMessage(`{}`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should just log and not panic
 	})
 
@@ -321,7 +322,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		select {
 		case <-checkCh:
 		case <-time.After(time.Second):
@@ -335,7 +336,7 @@ func TestHandleMessage(t *testing.T) {
 		credCh := make(chan *pb.CredentialTestResult, 1)
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: json.RawMessage(`not json`)}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should log error but not panic
 	})
 
@@ -346,7 +347,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should log error but not panic
 	})
 
@@ -357,7 +358,7 @@ func TestHandleMessage(t *testing.T) {
 		monCh := make(chan *pb.MonitoringCheck, 1)
 		checkCh := make(chan *pb.CheckResult, 1)
 		payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "check_jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		// Should log error but not panic
 	})
 
@@ -384,7 +385,7 @@ func TestHandleMessage(t *testing.T) {
 			JobType:    pb.JobType_DISCOVER,
 			SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
 		})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "discovery_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "discovery_job", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		select {
 		case <-snmpCh:
 		case <-time.After(500 * time.Millisecond):
@@ -412,7 +413,7 @@ func TestHandleMessage(t *testing.T) {
 			JobType:        pb.JobType_MIKROTIK,
 			MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", SshPort: 22, Username: "admin", Password: "pass"},
 		})
-		_, _ = handleMessage(context.Background(), channelMsg{Event: "backup_job", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+		_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "backup_job", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 		select {
 		case result := <-mtCh:
 			if result.Error != "" {
@@ -436,7 +437,7 @@ func TestHandleMessageRejectsOversizedPayload(t *testing.T) {
 	encoded := base64.StdEncoding.EncodeToString(oversized)
 	payload, _ := json.Marshal(map[string]string{"binary": encoded})
 
-	_, _ = handleMessage(context.Background(), channelMsg{Event: "jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
+	_, _ = handleMessage(context.Background(), channelMsg{Topic: "agent:test", Event: "jobs", Payload: payload}, "agent:test", testPools(t), snmpCh, mtCh, credCh, monCh, checkCh, make(chan *pb.LldpTopologyResult, 1))
 
 	// Verify no jobs were dispatched
 	select {
@@ -833,8 +834,9 @@ func TestDispatchJobCancelledContext(t *testing.T) {
 
 // fakeWSServer is a test helper that sets up a WebSocket server for runSession tests.
 type fakeWSServer struct {
-	ln   net.Listener
-	conn net.Conn
+	ln    net.Listener
+	conn  net.Conn
+	topic string
 }
 
 func newFakeWSServer(t *testing.T) *fakeWSServer {
@@ -867,13 +869,21 @@ func (s *fakeWSServer) acceptAndJoin(t *testing.T) {
 	resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
 	_, _ = conn.Write([]byte(resp))
 
-	// Read join frame
-	frameBuf := make([]byte, 4096)
-	_, _ = conn.Read(frameBuf)
+	// Read join frame and extract topic
+	joinPayload, err := readMaskedFrame(conn)
+	if err == nil {
+		var joinMsg channelMsg
+		if err := json.Unmarshal(joinPayload, &joinMsg); err == nil && joinMsg.Topic != "" {
+			s.topic = joinMsg.Topic
+		}
+	}
+	if s.topic == "" {
+		s.topic = "agent:agent-0"
+	}
 
 	// Send join OK
 	reply, _ := json.Marshal(channelMsg{
-		Topic:   "agent:agent-0",
+		Topic:   s.topic,
 		Event:   "phx_reply",
 		Payload: json.RawMessage(`{"status":"ok"}`),
 		Ref:     strPtr("1"),
@@ -884,7 +894,7 @@ func (s *fakeWSServer) acceptAndJoin(t *testing.T) {
 // sendEvent sends a channel message to the connected client.
 func (s *fakeWSServer) sendEvent(event string, payload json.RawMessage) {
 	msg, _ := json.Marshal(channelMsg{
-		Topic:   "agent:agent-0",
+		Topic:   s.topic,
 		Event:   event,
 		Payload: payload,
 	})

checks.go

@@ -29,10 +29,8 @@ var (
 		MaxIdleConnsPerHost: 10,
 		IdleConnTimeout:     90 * time.Second,
 	}
-	sslRootCAsOnce sync.Once
 	sslRootCAsMu   sync.Mutex
 	sslRootCAsPool *x509.CertPool
-	sslRootCAsErr  error
 
 	// sslRootCAs returns the system cert pool, cached after first successful load.
 	// Errors are not cached — subsequent calls will retry loading.
@@ -40,14 +38,15 @@ var (
 	sslRootCAs = func() (*x509.CertPool, error) {
 		sslRootCAsMu.Lock()
 		defer sslRootCAsMu.Unlock()
-		sslRootCAsOnce.Do(func() {
-			sslRootCAsPool, sslRootCAsErr = x509.SystemCertPool()
-		})
-		if sslRootCAsErr != nil {
-			// Reset the once so next call will retry
-			sslRootCAsOnce = sync.Once{}
+		if sslRootCAsPool != nil {
+			return sslRootCAsPool, nil
 		}
-		return sslRootCAsPool, sslRootCAsErr
+		pool, err := x509.SystemCertPool()
+		if err != nil {
+			return nil, err
+		}
+		sslRootCAsPool = pool
+		return pool, nil
 	}
 )
 
@@ -168,6 +167,7 @@ func executeHTTPCheck(ctx context.Context, config *pb.HttpCheckConfig, timeoutMs
 	if config.Regex != "" {
 		re, err := regexp.Compile(config.Regex)
 		if err != nil {
+			_, _ = io.Copy(io.Discard, resp.Body)
 			return 3, fmt.Sprintf("Invalid regex: %v", err), responseTime
 		}
 		body, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))

lldp.go

@@ -27,9 +27,15 @@ const (
 func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.LldpTopologyResult) {
 	deviceID := job.DeviceId
 	jobID := job.JobId
+	timestamp := time.Now().Unix()
 
 	if job.SnmpDevice == nil {
 		slog.Error("missing SNMP config for LLDP job", "job_id", jobID, "device_id", deviceID)
+		sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{
+			DeviceId:  deviceID,
+			JobId:     jobID,
+			Timestamp: timestamp,
+		}, jobID)
 		return
 	}
 
@@ -37,6 +43,11 @@ func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan
 	client, err := newSnmpConn(snmpDev)
 	if err != nil {
 		slog.Error("failed to connect SNMP for LLDP", "job_id", jobID, "device_id", deviceID, "error", err)
+		sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{
+			DeviceId:  deviceID,
+			JobId:     jobID,
+			Timestamp: timestamp,
+		}, jobID)
 		return
 	}
 	defer func() {
@@ -48,14 +59,25 @@ func executeLldpTopologyJob(ctx context.Context, job *pb.AgentJob, resultCh chan
 	result, err := discoverLldpNeighbors(client, deviceID, jobID)
 	if err != nil {
 		slog.Error("LLDP discovery failed", "job_id", jobID, "device_id", deviceID, "error", err)
+		sendLldpResultWithTimeout(ctx, resultCh, &pb.LldpTopologyResult{
+			DeviceId:  deviceID,
+			JobId:     jobID,
+			Timestamp: timestamp,
+		}, jobID)
 		return
 	}
 
+	sendLldpResultWithTimeout(ctx, resultCh, result, jobID)
+	slog.Info("LLDP topology discovered", "job_id", jobID, "device_id", deviceID, "neighbors", len(result.Neighbors))
+}
+
+func sendLldpResultWithTimeout(ctx context.Context, resultCh chan<- *pb.LldpTopologyResult, result *pb.LldpTopologyResult, jobID string) {
+	sendCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
 	select {
 	case resultCh <- result:
-		slog.Info("LLDP topology discovered", "job_id", jobID, "device_id", deviceID, "neighbors", len(result.Neighbors))
-	case <-ctx.Done():
-		slog.Warn("LLDP result send cancelled", "job_id", jobID)
+	case <-sendCtx.Done():
+		slog.Error("LLDP result send timeout - agent overloaded", "job_id", jobID)
 	}
 }
 

snmp.go

@@ -67,9 +67,11 @@ func executeSnmpJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.S
 	}
 	oidValues := make(map[string]string, totalOIDs)
 
+	cancelled := false
 	for _, q := range job.Queries {
 		if ctx.Err() != nil {
-			return
+			cancelled = true
+			break
 		}
 		switch q.QueryType {
 		case pb.QueryType_GET:
@@ -115,6 +117,10 @@ func executeSnmpJob(ctx context.Context, job *pb.AgentJob, resultCh chan<- *pb.S
 		}
 	}
 
+	if cancelled {
+		slog.Warn("snmp job cancelled, sending partial result", "job_id", job.JobId, "oids", len(oidValues))
+	}
+
 	result := &pb.SnmpResult{
 		DeviceId:  job.DeviceId,
 		JobType:   job.JobType,

websocket.go

@@ -28,11 +28,12 @@ func computeAcceptKey(key string) string {
 }
 
 const (
-	opText   = 1
-	opBinary = 2
-	opClose  = 8
-	opPing   = 9
-	opPong   = 10
+	opText       = 1
+	opBinary     = 2
+	opClose      = 8
+	opPing       = 9
+	opPong       = 10
+	opContinuation = 0
 
 	maxFrameSize = 16 << 20 // 16 MB
 )
@@ -219,22 +220,52 @@ func headerHasToken(value, token string) bool {
 }
 
 // ReadMessage reads the next text or binary message, handling control frames internally.
+// Fragmented messages (continuation frames per RFC 6455 Section 5.4) are reassembled.
 func (ws *WSConn) ReadMessage() ([]byte, int, error) {
+	var accumulated []byte
+	var msgOpcode int
+
 	for {
-		opcode, payload, err := ws.readFrame()
+		fin, opcode, payload, err := ws.readFrame()
 		if err != nil {
 			return nil, 0, err
 		}
+
 		switch opcode {
-		case opText, opBinary:
-			return payload, opcode, nil
 		case opPing:
 			if err := ws.writeFrame(opPong, payload); err != nil {
 				return nil, 0, fmt.Errorf("pong: %w", err)
 			}
+			continue
 		case opClose:
-			_ = ws.writeFrame(opClose, nil) // best-effort close reply
+			_ = ws.writeFrame(opClose, nil)
 			return nil, opClose, io.EOF
+		case opPong:
+			continue
+		}
+
+		if opcode == opText || opcode == opBinary {
+			// New message — if we already had an incomplete fragmented message, discard it
+			// (protocol violation by the server, but be resilient).
+			if !fin {
+				msgOpcode = opcode
+				accumulated = append(accumulated[:0], payload...)
+				continue
+			}
+			// Single unfragmented message
+			return payload, opcode, nil
+		}
+
+		// Continuation frame (opcode 0)
+		if accumulated == nil {
+			// Unexpected continuation without a start frame — skip
+			continue
+		}
+		accumulated = append(accumulated, payload...)
+		if fin {
+			result := accumulated
+			accumulated = nil
+			return result, msgOpcode, nil
 		}
 	}
 }
@@ -261,19 +292,20 @@ func (ws *WSConn) Close() error {
 const wsReadTimeout = 90 * time.Second
 const wsWriteTimeout = 30 * time.Second
 
-func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
+func (ws *WSConn) readFrame() (fin bool, opcode int, payload []byte, err error) {
 	// Set read deadline to detect dead connections (MEDIUM-4)
 	if nc, ok := ws.conn.(net.Conn); ok {
 		if err := nc.SetReadDeadline(time.Now().Add(wsReadTimeout)); err != nil {
-			return 0, nil, fmt.Errorf("set read deadline: %w", err)
+			return false, 0, nil, fmt.Errorf("set read deadline: %w", err)
 		}
 	}
 
 	var header [2]byte
 	if _, err = io.ReadFull(ws.reader, header[:]); err != nil {
-		return 0, nil, err
+		return false, 0, nil, err
 	}
 
+	fin = header[0]&0x80 != 0
 	opcode = int(header[0] & 0x0F)
 	masked := header[1]&0x80 != 0
 	length := uint64(header[1] & 0x7F)
@@ -282,32 +314,32 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
 	case 126:
 		var ext [2]byte
 		if _, err = io.ReadFull(ws.reader, ext[:]); err != nil {
-			return 0, nil, err
+			return false, 0, nil, err
 		}
 		length = uint64(binary.BigEndian.Uint16(ext[:]))
 	case 127:
 		var ext [8]byte
 		if _, err = io.ReadFull(ws.reader, ext[:]); err != nil {
-			return 0, nil, err
+			return false, 0, nil, err
 		}
 		length = binary.BigEndian.Uint64(ext[:])
 	}
 
 	if length > maxFrameSize {
-		return 0, nil, fmt.Errorf("frame size %d exceeds max %d", length, maxFrameSize)
+		return false, 0, nil, fmt.Errorf("frame size %d exceeds max %d", length, maxFrameSize)
 	}
 
 	var maskKey [4]byte
 	if masked {
 		if _, err = io.ReadFull(ws.reader, maskKey[:]); err != nil {
-			return 0, nil, err
+			return false, 0, nil, err
 		}
 	}
 
 	payload = make([]byte, length)
 	if length > 0 {
 		if _, err = io.ReadFull(ws.reader, payload); err != nil {
-			return 0, nil, err
+			return false, 0, nil, err
 		}
 	}
 
@@ -317,7 +349,7 @@ func (ws *WSConn) readFrame() (opcode int, payload []byte, err error) {
 		}
 	}
 
-	return opcode, payload, nil
+	return fin, opcode, payload, nil
 }
 
 func (ws *WSConn) writeFrame(opcode int, payload []byte) error {

websocket_test.go

@@ -120,7 +120,7 @@ func TestReadFrame(t *testing.T) {
 	buf.Write(payload)
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	opcode, data, err := ws.readFrame()
+	_, opcode, data, err := ws.readFrame()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -146,7 +146,7 @@ func TestReadFrame16BitLength(t *testing.T) {
 	buf.Write(payload)
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	opcode, data, err := ws.readFrame()
+	_, opcode, data, err := ws.readFrame()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -169,7 +169,7 @@ func TestReadFrame64BitLength(t *testing.T) {
 	buf.Write(payload)
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	opcode, data, err := ws.readFrame()
+	_, opcode, data, err := ws.readFrame()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -196,7 +196,7 @@ func TestReadFrameMasked(t *testing.T) {
 	buf.Write(masked)
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	opcode, data, err := ws.readFrame()
+	_, opcode, data, err := ws.readFrame()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -265,7 +265,7 @@ func TestReadFrame16BitLengthError(t *testing.T) {
 	buf.WriteByte(126)  // 16-bit length indicator, but no length bytes follow
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	_, _, err := ws.readFrame()
+	_, _, _, err := ws.readFrame()
 	if err == nil {
 		t.Error("expected error for truncated 16-bit length")
 	}
@@ -277,7 +277,7 @@ func TestReadFrame64BitLengthError(t *testing.T) {
 	buf.WriteByte(127)  // 64-bit length indicator, but no length bytes follow
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	_, _, err := ws.readFrame()
+	_, _, _, err := ws.readFrame()
 	if err == nil {
 		t.Error("expected error for truncated 64-bit length")
 	}
@@ -289,7 +289,7 @@ func TestReadFrameMaskKeyError(t *testing.T) {
 	buf.WriteByte(0x80 | 0x01) // masked + length 1, but no mask key or payload
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	_, _, err := ws.readFrame()
+	_, _, _, err := ws.readFrame()
 	if err == nil {
 		t.Error("expected error for missing mask key")
 	}
@@ -301,7 +301,7 @@ func TestReadFramePayloadError(t *testing.T) {
 	buf.WriteByte(5)    // length 5, but no payload
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	_, _, err := ws.readFrame()
+	_, _, _, err := ws.readFrame()
 	if err == nil {
 		t.Error("expected error for truncated payload")
 	}
@@ -348,7 +348,7 @@ func TestReadFrameExceedsMaxSize(t *testing.T) {
 	// Don't need to write the payload — should reject before reading it
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	_, _, err := ws.readFrame()
+	_, _, _, err := ws.readFrame()
 	if err == nil {
 		t.Error("expected error for frame exceeding max size")
 	}
@@ -363,7 +363,7 @@ func TestReadFrameEmptyPayload(t *testing.T) {
 	buf.WriteByte(0)    // zero length
 
 	ws := testWSConn(&nopCloser{readWriter: &buf})
-	opcode, data, err := ws.readFrame()
+	_, opcode, data, err := ws.readFrame()
 	if err != nil {
 		t.Fatal(err)
 	}