Merge pull request #154 from shelltime/feat/daemon-coding-heartbeat

feat(daemon): add coding heartbeat tracking with offline persistence

Author: AnnatarHe

cmd/daemon/main.go

@@ -105,6 +105,17 @@ func main() {
 		}
 	}
 
+	// Start heartbeat resync service if codeTracking is enabled
+	if cfg.CodeTracking != nil && cfg.CodeTracking.Enabled != nil && *cfg.CodeTracking.Enabled {
+		heartbeatResyncService := daemon.NewHeartbeatResyncService(cfg)
+		if err := heartbeatResyncService.Start(ctx); err != nil {
+			slog.Error("Failed to start heartbeat resync service", slog.Any("err", err))
+		} else {
+			slog.Info("Heartbeat resync service started")
+			defer heartbeatResyncService.Stop()
+		}
+	}
+
 	// Create processor instance
 	processor := daemon.NewSocketHandler(&cfg, pubsub)
 

daemon/handlers.go

@@ -17,16 +17,26 @@ func SocketTopicProccessor(messages <-chan *message.Message) {
 		if err := json.Unmarshal(msg.Payload, &socketMsg); err != nil {
 			slog.ErrorContext(ctx, "failed to parse socket message", slog.Any("err", err))
 			msg.Nack()
+			continue
 		}
 
-		if socketMsg.Type == SocketMessageTypeSync {
-			err := handlePubSubSync(ctx, socketMsg.Payload)
-			if err != nil {
-				slog.ErrorContext(ctx, "failed to parse socket message", slog.Any("err", err))
-				msg.Nack()
-			} else {
-				msg.Ack()
-			}
+		var err error
+		switch socketMsg.Type {
+		case SocketMessageTypeSync:
+			err = handlePubSubSync(ctx, socketMsg.Payload)
+		case SocketMessageTypeHeartbeat:
+			err = handlePubSubHeartbeat(ctx, socketMsg.Payload)
+		default:
+			slog.ErrorContext(ctx, "unknown socket message type", slog.String("type", string(socketMsg.Type)))
+			msg.Nack()
+			continue
+		}
+
+		if err != nil {
+			slog.ErrorContext(ctx, "failed to handle socket message", slog.Any("err", err), slog.String("type", string(socketMsg.Type)))
+			msg.Nack()
+		} else {
+			msg.Ack()
 		}
 	}
 }

daemon/handlers.heartbeat.go

@@ -0,0 +1,80 @@
+package daemon
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"os"
+
+	"github.com/malamtime/cli/model"
+)
+
+func handlePubSubHeartbeat(ctx context.Context, socketMsgPayload interface{}) error {
+	pb, err := json.Marshal(socketMsgPayload)
+	if err != nil {
+		slog.Error("Failed to marshal the heartbeat payload again for unmarshal", slog.Any("payload", socketMsgPayload))
+		return err
+	}
+
+	var heartbeatPayload model.HeartbeatPayload
+	err = json.Unmarshal(pb, &heartbeatPayload)
+	if err != nil {
+		slog.Error("Failed to parse heartbeat payload", slog.Any("payload", socketMsgPayload))
+		return err
+	}
+
+	if len(heartbeatPayload.Heartbeats) == 0 {
+		slog.Debug("Empty heartbeat payload, skipping")
+		return nil
+	}
+
+	cfg, err := stConfig.ReadConfigFile(ctx)
+	if err != nil {
+		slog.Error("Failed to read config file", slog.Any("err", err))
+		return err
+	}
+
+	// Try to send to server
+	err = model.SendHeartbeatsToServer(ctx, cfg, heartbeatPayload)
+	if err != nil {
+		slog.Warn("Failed to send heartbeats to server, saving to local file", slog.Any("err", err))
+		// On failure, save to local file
+		if saveErr := saveHeartbeatToFile(heartbeatPayload); saveErr != nil {
+			slog.Error("Failed to save heartbeat to local file", slog.Any("err", saveErr))
+			return saveErr
+		}
+		// Return nil because we saved the data locally - don't nack the message
+		return nil
+	}
+
+	slog.Info("Successfully sent heartbeats to server", slog.Int("count", len(heartbeatPayload.Heartbeats)))
+	return nil
+}
+
+// saveHeartbeatToFile appends a heartbeat payload as a single JSON line to the log file
+func saveHeartbeatToFile(payload model.HeartbeatPayload) error {
+	logFilePath := os.ExpandEnv(fmt.Sprintf("%s/%s", "$HOME", model.HEARTBEAT_LOG_FILE))
+
+	// Open file for appending, create if not exists
+	file, err := os.OpenFile(logFilePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
+	if err != nil {
+		return fmt.Errorf("failed to open heartbeat log file: %w", err)
+	}
+	defer file.Close()
+
+	// Marshal payload to JSON
+	data, err := json.Marshal(payload)
+	if err != nil {
+		return fmt.Errorf("failed to marshal heartbeat payload: %w", err)
+	}
+
+	// Write as single line with newline
+	_, err = file.Write(append(data, '\n'))
+	if err != nil {
+		return fmt.Errorf("failed to write heartbeat to file: %w", err)
+	}
+
+	slog.Debug("Saved heartbeat to local file", slog.String("path", logFilePath), slog.Int("count", len(payload.Heartbeats)))
+	return nil
+}

daemon/heartbeat_resync.go

@@ -0,0 +1,181 @@
+package daemon
+
+import (
+	"bufio"
+	"context"
+	"encoding/json"
+	"fmt"
+	"log/slog"
+	"os"
+	"sync"
+	"time"
+
+	"github.com/malamtime/cli/model"
+)
+
+const (
+	// HeartbeatResyncInterval is the interval for retrying failed heartbeats
+	HeartbeatResyncInterval = 30 * time.Minute
+)
+
+// HeartbeatResyncService handles periodic resync of failed heartbeats
+type HeartbeatResyncService struct {
+	config   model.ShellTimeConfig
+	ticker   *time.Ticker
+	stopChan chan struct{}
+	wg       sync.WaitGroup
+}
+
+// NewHeartbeatResyncService creates a new heartbeat resync service
+func NewHeartbeatResyncService(config model.ShellTimeConfig) *HeartbeatResyncService {
+	return &HeartbeatResyncService{
+		config:   config,
+		stopChan: make(chan struct{}),
+	}
+}
+
+// Start begins the periodic resync job
+func (s *HeartbeatResyncService) Start(ctx context.Context) error {
+	s.ticker = time.NewTicker(HeartbeatResyncInterval)
+	s.wg.Add(1)
+
+	go func() {
+		defer s.wg.Done()
+
+		// Run once at startup
+		s.resync(ctx)
+
+		for {
+			select {
+			case <-s.ticker.C:
+				s.resync(ctx)
+			case <-s.stopChan:
+				return
+			case <-ctx.Done():
+				return
+			}
+		}
+	}()
+
+	slog.Info("Heartbeat resync service started", slog.Duration("interval", HeartbeatResyncInterval))
+	return nil
+}
+
+// Stop stops the resync service
+func (s *HeartbeatResyncService) Stop() {
+	if s.ticker != nil {
+		s.ticker.Stop()
+	}
+	close(s.stopChan)
+	s.wg.Wait()
+	slog.Info("Heartbeat resync service stopped")
+}
+
+// resync reads failed heartbeats from the log file and attempts to send them
+func (s *HeartbeatResyncService) resync(ctx context.Context) {
+	logFilePath := os.ExpandEnv(fmt.Sprintf("%s/%s", "$HOME", model.HEARTBEAT_LOG_FILE))
+
+	// Check if file exists
+	if _, err := os.Stat(logFilePath); os.IsNotExist(err) {
+		slog.Debug("No heartbeat log file found, nothing to resync")
+		return
+	}
+
+	// Read the file
+	file, err := os.Open(logFilePath)
+	if err != nil {
+		slog.Error("Failed to open heartbeat log file for resync", slog.Any("err", err))
+		return
+	}
+
+	var lines []string
+	scanner := bufio.NewScanner(file)
+	for scanner.Scan() {
+		line := scanner.Text()
+		if line != "" {
+			lines = append(lines, line)
+		}
+	}
+	file.Close()
+
+	if err := scanner.Err(); err != nil {
+		slog.Error("Error reading heartbeat log file", slog.Any("err", err))
+		return
+	}
+
+	if len(lines) == 0 {
+		slog.Debug("No failed heartbeats to resync")
+		return
+	}
+
+	slog.Info("Starting heartbeat resync", slog.Int("pendingCount", len(lines)))
+
+	// Process each line
+	var failedLines []string
+	successCount := 0
+
+	for _, line := range lines {
+		var payload model.HeartbeatPayload
+		if err := json.Unmarshal([]byte(line), &payload); err != nil {
+			slog.Error("Failed to parse heartbeat line, discarding", slog.Any("err", err), slog.String("line", line))
+			continue
+		}
+
+		// Try to send to server
+		if err := model.SendHeartbeatsToServer(ctx, s.config, payload); err != nil {
+			slog.Warn("Failed to resync heartbeat, keeping for next retry", slog.Any("err", err))
+			failedLines = append(failedLines, line)
+		} else {
+			successCount++
+		}
+	}
+
+	// Rewrite the file with only failed lines
+	if err := s.rewriteLogFile(logFilePath, failedLines); err != nil {
+		slog.Error("Failed to update heartbeat log file", slog.Any("err", err))
+		return
+	}
+
+	slog.Info("Heartbeat resync completed",
+		slog.Int("success", successCount),
+		slog.Int("remaining", len(failedLines)))
+}
+
+// rewriteLogFile atomically rewrites the log file with the given lines
+func (s *HeartbeatResyncService) rewriteLogFile(logFilePath string, lines []string) error {
+	// If no lines remaining, remove the file
+	if len(lines) == 0 {
+		if err := os.Remove(logFilePath); err != nil && !os.IsNotExist(err) {
+			return fmt.Errorf("failed to remove empty log file: %w", err)
+		}
+		return nil
+	}
+
+	// Write to temp file first
+	tempFile := logFilePath + ".tmp"
+	file, err := os.OpenFile(tempFile, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
+	if err != nil {
+		return fmt.Errorf("failed to create temp file: %w", err)
+	}
+
+	for _, line := range lines {
+		if _, err := file.WriteString(line + "\n"); err != nil {
+			file.Close()
+			os.Remove(tempFile)
+			return fmt.Errorf("failed to write to temp file: %w", err)
+		}
+	}
+
+	if err := file.Close(); err != nil {
+		os.Remove(tempFile)
+		return fmt.Errorf("failed to close temp file: %w", err)
+	}
+
+	// Atomic rename
+	if err := os.Rename(tempFile, logFilePath); err != nil {
+		os.Remove(tempFile)
+		return fmt.Errorf("failed to rename temp file: %w", err)
+	}
+
+	return nil
+}

daemon/socket.go

@@ -14,7 +14,8 @@ import (
 type SocketMessageType string
 
 const (
-	SocketMessageTypeSync SocketMessageType = "sync"
+	SocketMessageTypeSync      SocketMessageType = "sync"
+	SocketMessageTypeHeartbeat SocketMessageType = "heartbeat"
 )
 
 type SocketMessage struct {
@@ -112,6 +113,22 @@ func (p *SocketHandler) handleConnection(conn net.Conn) {
 		if err := p.channel.Publish(PubSubTopic, chMsg); err != nil {
 			slog.Error("Error to publish topic", slog.Any("err", err))
 		}
+	case SocketMessageTypeHeartbeat:
+		// Only process heartbeat if codeTracking is enabled
+		if p.config.CodeTracking == nil || p.config.CodeTracking.Enabled == nil || !*p.config.CodeTracking.Enabled {
+			slog.Debug("Heartbeat message received but codeTracking is disabled, ignoring")
+			return
+		}
+		buf, err := json.Marshal(msg)
+		if err != nil {
+			slog.Error("Error encoding heartbeat message", slog.Any("err", err))
+			return
+		}
+
+		chMsg := message.NewMessage(watermill.NewUUID(), buf)
+		if err := p.channel.Publish(PubSubTopic, chMsg); err != nil {
+			slog.Error("Error publishing heartbeat topic", slog.Any("err", err))
+		}
 	default:
 		slog.Error("Unknown message type:", slog.String("messageType", string(msg.Type)))
 	}

model/api_heartbeat.go

@@ -0,0 +1,35 @@
+package model
+
+import (
+	"context"
+	"net/http"
+	"time"
+)
+
+// SendHeartbeatsToServer sends heartbeat data to the server
+func SendHeartbeatsToServer(ctx context.Context, cfg ShellTimeConfig, payload HeartbeatPayload) error {
+	ctx, span := modelTracer.Start(ctx, "api.sendHeartbeats")
+	defer span.End()
+
+	endpoint := Endpoint{
+		Token:       cfg.Token,
+		APIEndpoint: cfg.APIEndpoint,
+	}
+
+	var response HeartbeatResponse
+	err := SendHTTPRequestJSON(HTTPRequestOptions[HeartbeatPayload, HeartbeatResponse]{
+		Context:  ctx,
+		Endpoint: endpoint,
+		Method:   http.MethodPost,
+		Path:     "/api/v1/heartbeats",
+		Payload:  payload,
+		Response: &response,
+		Timeout:  10 * time.Second,
+	})
+
+	if err != nil {
+		return err
+	}
+
+	return nil
+}