fix: average prompt always 0 in WAKATIME_AI_STATS

Author: thanhhaudev

pkg/container/calculator.go

@@ -25,7 +25,10 @@ type LanguageStats struct {
 	Languages      map[string][2]interface{}
 }
 
-// AIStats stores AI vs human attribution aggregated across WakaTime projects.
+// AIStats stores AI vs human attribution from WakaTime's top-level totals.
+// AvgPromptLength comes from doc-defined ai_average_prompt_length when WakaTime
+// returns it; PromptLength is the raw total chars typed (ai_prompt_length) used
+// as fallback since the average field is currently missing in API responses.
 // HasData is false when nothing meaningful was reported, signalling writers to
 // hide the block entirely (avoids rendering a section full of zeros).
 type AIStats struct {
@@ -36,39 +39,28 @@ type AIStats struct {
 	AIInputTokens   int64
 	AIOutputTokens  int64
 	AvgPromptLength float64
+	PromptLength    int64
 	HasData         bool
 }
 
-// CalculateAIStats aggregates AI attribution fields across all WakaTime projects.
-// Avg prompt length is weighted by ai_input_tokens — projects without AI prompts
-// (zero tokens) contribute nothing, preventing them from skewing the average.
+// CalculateAIStats reads AI attribution from the top-level WakaTime stats.
+// We trust the API's own aggregation rather than re-summing per-project items.
 func (d *DataContainer) CalculateAIStats() *AIStats {
 	if d.Data.WakaTime == nil {
 		return &AIStats{}
 	}
 
-	var s AIStats
-	var weightedPromptSum float64
-	var promptWeight int64
-
-	for _, p := range d.Data.WakaTime.Data.Projects {
-		s.AIAdditions += p.AIAdditions
-		s.AIDeletions += p.AIDeletions
-		s.HumanAdditions += p.HumanAdditions
-		s.HumanDeletions += p.HumanDeletions
-		s.AIInputTokens += p.AIInputTokens
-		s.AIOutputTokens += p.AIOutputTokens
-
-		if p.AIInputTokens > 0 && p.AIAveragePromptLength > 0 {
-			weightedPromptSum += p.AIAveragePromptLength * float64(p.AIInputTokens)
-			promptWeight += p.AIInputTokens
-		}
-	}
-
-	if promptWeight > 0 {
-		s.AvgPromptLength = weightedPromptSum / float64(promptWeight)
+	src := d.Data.WakaTime.Data
+	s := AIStats{
+		AIAdditions:     src.AIAdditions,
+		AIDeletions:     src.AIDeletions,
+		HumanAdditions:  src.HumanAdditions,
+		HumanDeletions:  src.HumanDeletions,
+		AIInputTokens:   src.AIInputTokens,
+		AIOutputTokens:  src.AIOutputTokens,
+		AvgPromptLength: src.AIAvgPromptLength,
+		PromptLength:    src.AIPromptLength,
 	}
-
 	s.HasData = s.AIAdditions > 0 || s.AIInputTokens > 0
 	return &s
 }

pkg/container/calculator_test.go

@@ -104,65 +104,73 @@ func TestCalculateStreaks(t *testing.T) {
 	}
 }
 
-func newAIStatsContainer(projects []wakatime.StatsItem) *DataContainer {
+func newAIStatsContainer(stats *wakatime.Stats) *DataContainer {
 	d := NewDataContainer(log.Default(), &ClientManager{}, &config.Config{})
-	if projects != nil {
-		d.Data.WakaTime = &wakatime.Stats{}
-		d.Data.WakaTime.Data.Projects = projects
+	if stats != nil {
+		d.Data.WakaTime = stats
 	}
 	return d
 }
 
+func aiStats(setup func(s *wakatime.Stats)) *wakatime.Stats {
+	s := &wakatime.Stats{}
+	setup(s)
+	return s
+}
+
 func TestCalculateAIStats(t *testing.T) {
 	tests := []struct {
-		name     string
-		projects []wakatime.StatsItem
-		want     AIStats
+		name  string
+		stats *wakatime.Stats
+		want  AIStats
 	}{
 		{
-			name:     "no wakatime data",
-			projects: nil,
-			want:     AIStats{},
+			name:  "no wakatime data",
+			stats: nil,
+			want:  AIStats{},
 		},
 		{
-			name:     "wakatime present but no AI activity",
-			projects: []wakatime.StatsItem{{Name: "proj"}},
-			want:     AIStats{},
+			name:  "wakatime present but no AI activity",
+			stats: aiStats(func(s *wakatime.Stats) {}),
+			want:  AIStats{},
 		},
 		{
-			name: "aggregates additions and tokens across projects",
-			projects: []wakatime.StatsItem{
-				{Name: "alpha", AIAdditions: 100, HumanAdditions: 50, AIInputTokens: 1000, AIOutputTokens: 2000, AIAveragePromptLength: 100},
-				{Name: "beta", AIAdditions: 200, HumanAdditions: 150, AIInputTokens: 3000, AIOutputTokens: 4000, AIAveragePromptLength: 200},
-			},
+			name: "reads top-level totals including doc-compliant avg prompt",
+			stats: aiStats(func(s *wakatime.Stats) {
+				s.Data.AIAdditions = 300
+				s.Data.HumanAdditions = 200
+				s.Data.AIInputTokens = 4000
+				s.Data.AIOutputTokens = 6000
+				s.Data.AIAvgPromptLength = 175
+			}),
 			want: AIStats{
 				AIAdditions:     300,
 				HumanAdditions:  200,
 				AIInputTokens:   4000,
 				AIOutputTokens:  6000,
-				AvgPromptLength: (100*1000 + 200*3000) / 4000.0, // 175
+				AvgPromptLength: 175,
 				HasData:         true,
 			},
 		},
 		{
-			name: "skips projects with zero ai_input_tokens for prompt-length avg",
-			projects: []wakatime.StatsItem{
-				{Name: "with-ai", AIAdditions: 10, AIInputTokens: 500, AIAveragePromptLength: 80},
-				{Name: "no-ai", HumanAdditions: 100, AIAveragePromptLength: 9999}, // bogus, should be ignored
-			},
+			name: "exposes raw ai_prompt_length when avg field is missing",
+			stats: aiStats(func(s *wakatime.Stats) {
+				s.Data.AIAdditions = 10
+				s.Data.AIInputTokens = 500
+				s.Data.AIPromptLength = 484803
+			}),
 			want: AIStats{
-				AIAdditions:     10,
-				HumanAdditions:  100,
-				AIInputTokens:   500,
-				AvgPromptLength: 80,
-				HasData:         true,
+				AIAdditions:   10,
+				AIInputTokens: 500,
+				PromptLength:  484803,
+				HasData:       true,
 			},
 		},
 		{
 			name: "AI additions without tokens still triggers HasData",
-			projects: []wakatime.StatsItem{
-				{Name: "p", AIAdditions: 5},
-			},
+			stats: aiStats(func(s *wakatime.Stats) {
+				s.Data.AIAdditions = 5
+			}),
 			want: AIStats{
 				AIAdditions: 5,
 				HasData:     true,
@@ -172,7 +180,7 @@ func TestCalculateAIStats(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			d := newAIStatsContainer(tt.projects)
+			d := newAIStatsContainer(tt.stats)
 			got := d.CalculateAIStats()
 
 			if got.AIAdditions != tt.want.AIAdditions {
@@ -190,6 +198,9 @@ func TestCalculateAIStats(t *testing.T) {
 			if math.Abs(got.AvgPromptLength-tt.want.AvgPromptLength) > 0.001 {
 				t.Errorf("AvgPromptLength: got %v, want %v", got.AvgPromptLength, tt.want.AvgPromptLength)
 			}
+			if got.PromptLength != tt.want.PromptLength {
+				t.Errorf("PromptLength: got %d, want %d", got.PromptLength, tt.want.PromptLength)
+			}
 			if got.HasData != tt.want.HasData {
 				t.Errorf("HasData: got %v, want %v", got.HasData, tt.want.HasData)
 			}

pkg/container/container.go

@@ -38,7 +38,7 @@ func (d *DataContainer) metrics(com *CommitStats, lang *LanguageStats, ai *AISta
 	version := d.Config.ProgressBarVersion
 	aiBlock := ""
 	if ai != nil && ai.HasData {
-		aiBlock = writer.MakeAIStatsList(ai.AIAdditions, ai.HumanAdditions, ai.AIInputTokens, ai.AIOutputTokens, ai.AvgPromptLength)
+		aiBlock = writer.MakeAIStatsList(ai.AIAdditions, ai.HumanAdditions, ai.AIInputTokens, ai.AIOutputTokens, ai.PromptLength, ai.AvgPromptLength)
 	}
 	return map[string]string{
 		"LANGUAGE_PER_REPO":   writer.MakeLanguagePerRepoList(d.Data.Repositories, version),

pkg/wakatime/stats.go

@@ -21,15 +21,6 @@ type StatsItem struct {
 	Hours   int     `json:"hours"`
 	Minutes int     `json:"minutes"`
 	Seconds int     `json:"seconds"`
-
-	// AI attribution fields (populated on projects when the user has GenAI tooling integrated)
-	AIAdditions           int64   `json:"ai_additions"`
-	AIDeletions           int64   `json:"ai_deletions"`
-	HumanAdditions        int64   `json:"human_additions"`
-	HumanDeletions        int64   `json:"human_deletions"`
-	AIInputTokens         int64   `json:"ai_input_tokens"`
-	AIOutputTokens        int64   `json:"ai_output_tokens"`
-	AIAveragePromptLength float64 `json:"ai_average_prompt_length"`
 }
 
 type Stats struct {
@@ -40,6 +31,24 @@ type Stats struct {
 		Editors          []StatsItem `json:"editors"`
 		Projects         []StatsItem `json:"projects"`
 		OperatingSystems []StatsItem `json:"operating_systems"`
+
+		// AI attribution aggregates (top-level totals across the user's activity in this range).
+		AIAdditions    int64 `json:"ai_additions"`
+		AIDeletions    int64 `json:"ai_deletions"`
+		HumanAdditions int64 `json:"human_additions"`
+		HumanDeletions int64 `json:"human_deletions"`
+		AIInputTokens  int64 `json:"ai_input_tokens"`
+		AIOutputTokens int64 `json:"ai_output_tokens"`
+
+		// AIAvgPromptLength is doc-defined (avg chars per prompt). Preferred when populated.
+		AIAvgPromptLength float64 `json:"ai_average_prompt_length"`
+		// TODO(wakatime-api): remove AIPromptLength once WakaTime's /stats endpoint
+		// actually returns ai_average_prompt_length at top-level (currently doc-only).
+		// AIPromptLength is the raw total chars typed to AI tools and is semantically
+		// NOT an average — we render it under the "Average Prompt" label as a stop-gap
+		// per user decision. Drop this field + the writer fallback when the API is fixed.
+		// See: https://wakatime.com/developers#stats
+		AIPromptLength int64 `json:"ai_prompt_length"`
 	} `json:"data"`
 }
 

pkg/writer/writer.go

@@ -178,7 +178,11 @@ func MakeCodingStreakList(s *wakatime.AllTimeSinceTodayStats, currentStreak, lon
 
 // MakeAIStatsList returns a summary of AI vs human coding attribution from WakaTime.
 // Returns an empty string when there is no AI activity to display, so the block is hidden entirely.
-func MakeAIStatsList(aiAdd, humanAdd, inTokens, outTokens int64, avgPrompt float64) string {
+//
+// The "Average Prompt" row prefers avgPrompt (doc-defined ai_average_prompt_length) and
+// falls back to promptLength (ai_prompt_length) while the avg field is missing in API
+// responses. The row is omitted when neither value is populated.
+func MakeAIStatsList(aiAdd, humanAdd, inTokens, outTokens, promptLength int64, avgPrompt float64) string {
 	if aiAdd == 0 && inTokens == 0 {
 		return ""
 	}
@@ -189,14 +193,29 @@ func MakeAIStatsList(aiAdd, humanAdd, inTokens, outTokens int64, avgPrompt float
 		aiContribution = float64(aiAdd) / float64(totalAdd) * 100
 	}
 
+	// TODO(wakatime-api): drop the promptLength fallback once WakaTime ships
+	// ai_average_prompt_length at top-level. Today's API returns only ai_prompt_length
+	// (a raw total, not an average), so we display the total under the "Average Prompt"
+	// label as a stop-gap. When avgPrompt becomes non-zero, this branch becomes unreachable
+	// in practice and the whole switch can collapse to a single avgPrompt assignment.
+	var promptValue int64
+	switch {
+	case avgPrompt > 0:
+		promptValue = int64(math.Round(avgPrompt))
+	case promptLength > 0:
+		promptValue = promptLength
+	}
+
 	var b strings.Builder
 	b.WriteString("**🤖 My AI Footprint**\n\n")
 	b.WriteString("```text\n")
 	fmt.Fprintf(&b, "🤖 Generated by AI:        %s lines\n", addCommas(int(aiAdd)))
 	fmt.Fprintf(&b, "👤 Written by Hand:        %s lines\n", addCommas(int(humanAdd)))
 	fmt.Fprintf(&b, "📊 AI Contribution:        %.1f%%\n", aiContribution)
 	fmt.Fprintf(&b, "🔤 Tokens In / Out:        %s / %s\n", humanizeCount(inTokens), humanizeCount(outTokens))
-	fmt.Fprintf(&b, "💬 Average Prompt:         %d chars\n", int(math.Round(avgPrompt)))
+	if promptValue > 0 {
+		fmt.Fprintf(&b, "💬 Average Prompt:         %s chars\n", humanizeCount(promptValue))
+	}
 	b.WriteString("```\n\n")
 
 	return b.String()

pkg/writer/writer_test.go

@@ -122,60 +122,63 @@ func TestMakeCodingStreakList(t *testing.T) {
 
 func TestMakeAIStatsList(t *testing.T) {
 	tests := []struct {
-		name      string
-		aiAdd     int64
-		humanAdd  int64
-		inTokens  int64
-		outTokens int64
-		avgPrompt float64
-		empty     bool
-		contains  []string
+		name         string
+		aiAdd        int64
+		humanAdd     int64
+		inTokens     int64
+		outTokens    int64
+		promptLength int64
+		avgPrompt    float64
+		empty        bool
+		contains     []string
+		notContains  []string
 	}{
 		{
 			name:  "no activity returns empty",
 			empty: true,
 		},
 		{
-			name:      "only AI activity reported",
+			name:      "Average Prompt uses ai_average_prompt_length when populated",
 			aiAdd:     1500,
 			inTokens:  120_000,
 			outTokens: 350_000,
-			avgPrompt: 142,
+			avgPrompt: 7200,
 			contains: []string{
 				"🤖 My AI Footprint",
 				"Generated by AI:        1,500 lines",
 				"Written by Hand:        0 lines",
 				"AI Contribution:        100.0%",
 				"Tokens In / Out:        120.0K / 350.0K",
-				"Average Prompt:         142 chars",
+				"Average Prompt:         7.2K chars",
 			},
 		},
 		{
-			name:      "mixed AI and human contribution",
-			aiAdd:     12340,
-			humanAdd:  8721,
-			inTokens:  1_200_000,
-			outTokens: 3_400_000,
-			avgPrompt: 142.4,
+			name:         "Average Prompt falls back to ai_prompt_length when avg is zero",
+			aiAdd:        12340,
+			humanAdd:     8721,
+			inTokens:     1_200_000,
+			outTokens:    3_400_000,
+			promptLength: 484803,
 			contains: []string{
 				"AI Contribution:        58.6%",
 				"Tokens In / Out:        1.2M / 3.4M",
-				"Average Prompt:         142 chars",
+				"Average Prompt:         484.8K chars",
 			},
 		},
 		{
-			name:     "small token counts use addCommas",
+			name:     "small token counts use addCommas; prompt row hidden when both zero",
 			aiAdd:    20,
 			inTokens: 999,
 			contains: []string{
 				"Tokens In / Out:        999 / 0",
 			},
+			notContains: []string{"Average Prompt:"},
 		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			got := MakeAIStatsList(tt.aiAdd, tt.humanAdd, tt.inTokens, tt.outTokens, tt.avgPrompt)
+			got := MakeAIStatsList(tt.aiAdd, tt.humanAdd, tt.inTokens, tt.outTokens, tt.promptLength, tt.avgPrompt)
 
 			if tt.empty {
 				if got != "" {
@@ -189,6 +192,11 @@ func TestMakeAIStatsList(t *testing.T) {
 					t.Errorf("expected output to contain %q, got:\n%s", want, got)
 				}
 			}
+			for _, banned := range tt.notContains {
+				if strings.Contains(got, banned) {
+					t.Errorf("expected output to NOT contain %q, got:\n%s", banned, got)
+				}
+			}
 		})
 	}
 }