From 363a2c45ba5b6f8eb9ce45b282e4cb81370caccf Mon Sep 17 00:00:00 2001
From: Steve Yegge <steve.yegge@gmail.com>
Date: Sat, 20 Dec 2025 08:21:39 -0800
Subject: [PATCH] fix(statusline): Count rigs from all agent types, not just
 polecats
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Rig count now includes rigs with any active sessions (witness, refinery, crew)
- Previously only counted rigs that had active polecat sessions
- Updated tests to use categorizeSession instead of removed helper functions

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
---
 internal/cmd/statusline.go      | 10 ++++-
 internal/cmd/statusline_test.go | 69 ++++++++++++++++++---------------
 2 files changed, 46 insertions(+), 33 deletions(-)

diff --git a/internal/cmd/statusline.go b/internal/cmd/statusline.go
index f50e5f89..8856d6fc 100644
--- a/internal/cmd/statusline.go
+++ b/internal/cmd/statusline.go
@@ -93,7 +93,9 @@ func runMayorStatusLine(t *tmux.Tmux) error {
 		return nil // Silent fail
 	}
 
-	// Count only actual polecats (not witnesses, refineries, deacon, crew)
+	// Count polecats and rigs
+	// Polecats: only actual polecats (not witnesses, refineries, deacon, crew)
+	// Rigs: any rig with active sessions (witness, refinery, crew, or polecat)
 	polecatCount := 0
 	rigs := make(map[string]bool)
 	for _, s := range sessions {
@@ -101,9 +103,13 @@ func runMayorStatusLine(t *tmux.Tmux) error {
 		if agent == nil {
 			continue
 		}
+		// Count rigs from any rig-level agent (has non-empty Rig field)
+		if agent.Rig != "" {
+			rigs[agent.Rig] = true
+		}
+		// Count only polecats for polecat count
 		if agent.Type == AgentPolecat {
 			polecatCount++
-			rigs[agent.Rig] = true
 		}
 	}
 	rigCount := len(rigs)
diff --git a/internal/cmd/statusline_test.go b/internal/cmd/statusline_test.go
index 15975a0d..3d5ff915 100644
--- a/internal/cmd/statusline_test.go
+++ b/internal/cmd/statusline_test.go
@@ -2,10 +2,10 @@ package cmd
 
 import "testing"
 
-func TestExtractRigFromSession(t *testing.T) {
+func TestCategorizeSessionRig(t *testing.T) {
 	tests := []struct {
 		session string
-		want    string
+		wantRig string
 	}{
 		// Standard polecat sessions
 		{"gt-gastown-slit", "gastown"},
@@ -16,59 +16,66 @@ func TestExtractRigFromSession(t *testing.T) {
 		{"gt-gastown-crew-max", "gastown"},
 		{"gt-myrig-crew-user", "myrig"},
 
-		// Witness sessions (daemon.go style: gt-<rig>-witness)
+		// Witness sessions
 		{"gt-gastown-witness", "gastown"},
 		{"gt-myrig-witness", "myrig"},
 
-		// Witness sessions (witness.go style: gt-witness-<rig>)
-		{"gt-witness-gastown", "gastown"},
-		{"gt-witness-myrig", "myrig"},
-
 		// Refinery sessions
 		{"gt-gastown-refinery", "gastown"},
 		{"gt-myrig-refinery", "myrig"},
 
 		// Edge cases
-		{"gt-a-b", "a"},       // minimum valid
-		{"gt-ab", ""},         // too short, no worker
-		{"gt-", ""},           // invalid
-		{"gt", ""},            // invalid
+		{"gt-a-b", "a"}, // minimum valid
+
+		// Town-level agents (no rig)
+		{"gt-mayor", ""},
+		{"gt-deacon", ""},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.session, func(t *testing.T) {
-			got := extractRigFromSession(tt.session)
-			if got != tt.want {
-				t.Errorf("extractRigFromSession(%q) = %q, want %q", tt.session, got, tt.want)
+			agent := categorizeSession(tt.session)
+			gotRig := ""
+			if agent != nil {
+				gotRig = agent.Rig
+			}
+			if gotRig != tt.wantRig {
+				t.Errorf("categorizeSession(%q).Rig = %q, want %q", tt.session, gotRig, tt.wantRig)
 			}
 		})
 	}
 }
 
-func TestIsPolecatSession(t *testing.T) {
+func TestCategorizeSessionType(t *testing.T) {
 	tests := []struct {
-		session string
-		want    bool
+		session  string
+		wantType AgentType
 	}{
-		// Polecat sessions (should return true)
-		{"gt-gastown-slit", true},
-		{"gt-gastown-Toast", true},
-		{"gt-myrig-worker", true},
-		{"gt-a-b", true},
+		// Polecat sessions
+		{"gt-gastown-slit", AgentPolecat},
+		{"gt-gastown-Toast", AgentPolecat},
+		{"gt-myrig-worker", AgentPolecat},
+		{"gt-a-b", AgentPolecat},
 
-		// Non-polecat sessions (should return false)
-		{"gt-gastown-witness", false},
-		{"gt-witness-gastown", false},
-		{"gt-gastown-refinery", false},
-		{"gt-gastown-crew-max", false},
-		{"gt-myrig-crew-user", false},
+		// Non-polecat sessions
+		{"gt-gastown-witness", AgentWitness},
+		{"gt-gastown-refinery", AgentRefinery},
+		{"gt-gastown-crew-max", AgentCrew},
+		{"gt-myrig-crew-user", AgentCrew},
+
+		// Town-level agents
+		{"gt-mayor", AgentMayor},
+		{"gt-deacon", AgentDeacon},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.session, func(t *testing.T) {
-			got := isPolecatSession(tt.session)
-			if got != tt.want {
-				t.Errorf("isPolecatSession(%q) = %v, want %v", tt.session, got, tt.want)
+			agent := categorizeSession(tt.session)
+			if agent == nil {
+				t.Fatalf("categorizeSession(%q) returned nil", tt.session)
+			}
+			if agent.Type != tt.wantType {
+				t.Errorf("categorizeSession(%q).Type = %v, want %v", tt.session, agent.Type, tt.wantType)
 			}
 		})
 	}