Merge pull request #1332 from wakatime/bugfix/offline-ai-heartbeats

Fix saving AI heartbeats to offline db

Author: alanhamlett

cmd/heartbeat/ai_sync_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/gandarez/go-realpath"
 	cmdheartbeat "github.com/wakatime/wakatime-cli/cmd/heartbeat"
 	"github.com/wakatime/wakatime-cli/pkg/ini"
+	"github.com/wakatime/wakatime-cli/pkg/offline"
 
 	"github.com/spf13/viper"
 	"github.com/stretchr/testify/assert"
@@ -210,3 +211,64 @@ func TestRunAISyncActivity_UsesAlternateProject(t *testing.T) {
 	assert.Equal(t, 0, code)
 	assert.Equal(t, 1, numCalls)
 }
+
+func TestRunAISyncActivity_RateLimitedWithoutEntity_SavesOffline(t *testing.T) {
+	resetSingleton(t)
+
+	home := t.TempDir()
+	t.Setenv("HOME", home)
+	t.Setenv("USERPROFILE", home)
+
+	tmpDir := t.TempDir()
+
+	tmpDir, err := realpath.Realpath(tmpDir)
+	require.NoError(t, err)
+
+	copyFile(t, "testdata/main.go", filepath.Join(tmpDir, "main.go"))
+
+	entity, err := filepath.Abs(filepath.Join(tmpDir, "main.go"))
+	require.NoError(t, err)
+
+	entity = strings.ReplaceAll(entity, "\\", "/")
+
+	transcriptDir := filepath.Join(home, ".claude", "projects", "sample-project")
+	require.NoError(t, os.MkdirAll(transcriptDir, 0o755))
+
+	transcriptPath := filepath.Join(transcriptDir, "session.jsonl")
+	transcript := strings.Join([]string{
+		"{\"timestamp\":\"2026-03-18T12:00:00Z\",\"version\":\"2.1.45\"," +
+			"\"toolUseResult\":{\"filePath\":\"" + entity + "\"," +
+			"\"structuredPatch\":[{\"oldLines\":3,\"newLines\":5},{\"oldLines\":4,\"newLines\":1}]}}",
+	}, "\n") + "\n"
+	require.NoError(t, os.WriteFile(transcriptPath, []byte(transcript), 0o644))
+
+	tmpInternalFile, err := os.CreateTemp(t.TempDir(), "wakatime-internal-config")
+	require.NoError(t, err)
+
+	defer tmpInternalFile.Close()
+
+	offlineQueueFile, err := os.CreateTemp(t.TempDir(), "offline-queue-file")
+	require.NoError(t, err)
+
+	defer offlineQueueFile.Close()
+
+	v := viper.New()
+	v.Set("heartbeat-rate-limit-seconds", 120)
+	v.Set("internal-config", tmpInternalFile.Name())
+	v.Set("internal.heartbeats_last_sent_at", time.Now().Format(ini.DateFormat))
+	v.Set("internal.ai_heartbeats_last_parsed_at", time.Date(2026, 3, 18, 11, 0, 0, 0, time.UTC).Format(ini.DateFormat))
+	v.Set("key", "00000000-0000-4000-8000-000000000000")
+	v.Set("offline-queue-file", offlineQueueFile.Name())
+	v.Set("plugin", "plugin/0.0.1")
+	v.Set("project", "myproject")
+	v.Set("project-folder", "/path/to/project")
+	v.Set("timeout", 5)
+
+	code, err := cmdheartbeat.RunAISyncActivity(t.Context(), v)
+	require.NoError(t, err)
+	assert.Equal(t, 0, code)
+
+	offlineCount, err := offline.CountHeartbeats(t.Context(), offlineQueueFile.Name())
+	require.NoError(t, err)
+	assert.Equal(t, 1, offlineCount)
+}

cmd/heartbeat/heartbeat.go

@@ -267,7 +267,7 @@ func sendPreparedHeartbeats(
 		LastSentAt: params.Offline.LastSentAt,
 		Timeout:    params.Offline.RateLimit,
 	}) {
-		err := offlinecmd.SaveHeartbeats(ctx, v, queueFilepath, heartbeats)
+		err := offlinecmd.SaveHeartbeatsWithParams(ctx, v, queueFilepath, heartbeats, params)
 		if err == nil {
 			return nil
 		}
@@ -288,7 +288,7 @@ func sendPreparedHeartbeats(
 		logger.Debugf("save %d extra heartbeat(s) to offline queue", len(extraHeartbeats))
 
 		go func(done chan<- bool) {
-			if err := offlinecmd.SaveHeartbeats(ctx, v, queueFilepath, extraHeartbeats); err != nil {
+			if err := offlinecmd.SaveHeartbeatsWithParams(ctx, v, queueFilepath, extraHeartbeats, params); err != nil {
 				logger.Errorf("failed to save extra heartbeats to offline queue: %s", err)
 			}
 
@@ -300,7 +300,7 @@ func sendPreparedHeartbeats(
 
 	sender, err := buildHandle(ctx, v, params, queueFilepath)
 	if err != nil {
-		if err := offlinecmd.SaveHeartbeats(ctx, v, queueFilepath, heartbeats); err != nil {
+		if err := offlinecmd.SaveHeartbeatsWithParams(ctx, v, queueFilepath, heartbeats, params); err != nil {
 			logger.Errorf("failed to save extra heartbeats to offline queue: %s", err)
 		}
 

cmd/offline/offline.go

@@ -30,6 +30,27 @@ func SaveHeartbeats(
 		return fmt.Errorf("failed to load command parameters: %w", err)
 	}
 
+	return saveHeartbeats(ctx, v, queueFilepath, heartbeats, params)
+}
+
+// SaveHeartbeatsWithParams saves heartbeats to the offline db using already loaded command params.
+func SaveHeartbeatsWithParams(
+	ctx context.Context,
+	v *viper.Viper,
+	queueFilepath string,
+	heartbeats []heartbeat.Heartbeat,
+	params params.Params,
+) error {
+	return saveHeartbeats(ctx, v, queueFilepath, heartbeats, params)
+}
+
+func saveHeartbeats(
+	ctx context.Context,
+	v *viper.Viper,
+	queueFilepath string,
+	heartbeats []heartbeat.Heartbeat,
+	params params.Params,
+) error {
 	logger := log.Extract(ctx)
 
 	setLogFields(ctx, params)