Move token usage calculation into agent logic

cmd/entire/cli/agent/agent.go

@@ -11,8 +11,12 @@ import (
 // Each agent implementation (Claude Code, Cursor, Aider, etc.) converts its
 // native format to the normalized types defined in this package.
 type Agent interface {
-	// Name returns the agent identifier (e.g., "claude-code", "cursor")
-	Name() string
+	// Name returns the agent registry key (e.g., "claude-code", "gemini")
+	Name() AgentName
+
+	// Type returns the agent type identifier (e.g., "Claude Code", "Gemini CLI")
+	// This is stored in metadata and trailers.
+	Type() AgentType
 
 	// Description returns a human-readable description for UI
 	Description() string

cmd/entire/cli/agent/agent_test.go

@@ -6,14 +6,16 @@ import (
 	"time"
 )
 
-const mockAgentName = "mock" // Used by mock implementations
+const mockAgentName AgentName = "mock" // Used by mock implementations
+const mockAgentType AgentType = "Mock Agent"
 
 // mockAgent is a minimal implementation of Agent for testing interface compliance.
 type mockAgent struct{}
 
 var _ Agent = (*mockAgent)(nil) // Compile-time interface check
 
-func (m *mockAgent) Name() string                  { return mockAgentName }
+func (m *mockAgent) Name() AgentName               { return mockAgentName }
+func (m *mockAgent) Type() AgentType               { return mockAgentType }
 func (m *mockAgent) Description() string           { return "Mock agent for testing" }
 func (m *mockAgent) DetectPresence() (bool, error) { return false, nil }
 func (m *mockAgent) GetHookConfigPath() string     { return "" }

cmd/entire/cli/agent/chunking.go

@@ -19,7 +19,7 @@ const (
 // ChunkTranscript splits a transcript into chunks using the appropriate agent.
 // If agentType is empty or the agent doesn't implement TranscriptChunker,
 // falls back to JSONL (line-based) chunking.
-func ChunkTranscript(content []byte, agentType string) ([][]byte, error) {
+func ChunkTranscript(content []byte, agentType AgentType) ([][]byte, error) {
 	if len(content) <= MaxChunkSize {
 		return [][]byte{content}, nil
 	}
@@ -45,7 +45,7 @@ func ChunkTranscript(content []byte, agentType string) ([][]byte, error) {
 // ReassembleTranscript combines chunks back into a single transcript.
 // If agentType is empty or the agent doesn't implement TranscriptChunker,
 // falls back to JSONL (line-based) reassembly.
-func ReassembleTranscript(chunks [][]byte, agentType string) ([]byte, error) {
+func ReassembleTranscript(chunks [][]byte, agentType AgentType) ([]byte, error) {
 	if len(chunks) == 0 {
 		return nil, nil
 	}
@@ -169,9 +169,9 @@ type geminiTranscriptDetect struct {
 }
 
 // DetectAgentTypeFromContent detects the agent type from transcript content.
-// Returns "Gemini CLI" if it appears to be Gemini JSON format, empty string otherwise.
+// Returns AgentTypeGemini if it appears to be Gemini JSON format, empty AgentType otherwise.
 // This is used when the agent type is unknown but we need to chunk/reassemble correctly.
-func DetectAgentTypeFromContent(content []byte) string {
+func DetectAgentTypeFromContent(content []byte) AgentType {
 	// Quick check: Gemini JSON starts with { and has a messages array
 	trimmed := strings.TrimSpace(string(content))
 	if !strings.HasPrefix(trimmed, "{") {
@@ -186,7 +186,7 @@ func DetectAgentTypeFromContent(content []byte) string {
 
 	// Must have at least one message to be considered Gemini format
 	if len(transcript.Messages) > 0 {
-		return "Gemini CLI"
+		return AgentTypeGemini
 	}
 
 	return ""

cmd/entire/cli/agent/chunking_test.go

@@ -197,12 +197,12 @@ func TestDetectAgentTypeFromContent(t *testing.T) {
 	tests := []struct {
 		name     string
 		content  []byte
-		expected string
+		expected AgentType
 	}{
 		{
 			name:     "Gemini JSON",
 			content:  []byte(`{"messages":[{"type":"user","content":"hi"}]}`),
-			expected: "Gemini CLI",
+			expected: AgentTypeGemini,
 		},
 		{
 			name:     "JSONL",

cmd/entire/cli/agent/claudecode/claude.go

@@ -31,11 +31,16 @@ func NewClaudeCodeAgent() agent.Agent {
 	return &ClaudeCodeAgent{}
 }
 
-// Name returns the agent identifier.
-func (c *ClaudeCodeAgent) Name() string {
+// Name returns the agent registry key.
+func (c *ClaudeCodeAgent) Name() agent.AgentName {
 	return agent.AgentNameClaudeCode
 }
 
+// Type returns the agent type identifier.
+func (c *ClaudeCodeAgent) Type() agent.AgentType {
+	return agent.AgentTypeClaudeCode
+}
+
 // Description returns a human-readable description.
 func (c *ClaudeCodeAgent) Description() string {
 	return "Claude Code - Anthropic's CLI coding assistant"

cmd/entire/cli/agent/claudecode/transcript.go

@@ -9,7 +9,7 @@ import (
 	"path/filepath"
 	"strings"
 
-	"entire.io/cli/cmd/entire/cli/checkpoint"
+	"entire.io/cli/cmd/entire/cli/agent"
 )
 
 // Transcript parsing types - Claude Code uses JSONL format
@@ -227,7 +227,7 @@ func FindCheckpointUUID(lines []TranscriptLine, toolUseID string) (string, bool)
 //
 // Due to streaming, multiple transcript rows may share the same message.id.
 // We deduplicate by taking the row with the highest output_tokens for each message.id.
-func CalculateTokenUsage(transcript []TranscriptLine) *checkpoint.TokenUsage {
+func CalculateTokenUsage(transcript []TranscriptLine) *agent.TokenUsage {
 	// Map from message.id to the usage with highest output_tokens
 	usageByMessageID := make(map[string]messageUsage)
 
@@ -253,7 +253,7 @@ func CalculateTokenUsage(transcript []TranscriptLine) *checkpoint.TokenUsage {
 	}
 
 	// Sum up all unique messages
-	usage := &checkpoint.TokenUsage{
+	usage := &agent.TokenUsage{
 		APICallCount: len(usageByMessageID),
 	}
 	for _, u := range usageByMessageID {
@@ -268,9 +268,9 @@ func CalculateTokenUsage(transcript []TranscriptLine) *checkpoint.TokenUsage {
 
 // CalculateTokenUsageFromFile calculates token usage from a Claude Code transcript file.
 // If startLine > 0, only considers lines from startLine onwards.
-func CalculateTokenUsageFromFile(path string, startLine int) (*checkpoint.TokenUsage, error) {
+func CalculateTokenUsageFromFile(path string, startLine int) (*agent.TokenUsage, error) {
 	if path == "" {
-		return &checkpoint.TokenUsage{}, nil
+		return &agent.TokenUsage{}, nil
 	}
 
 	transcript, err := parseTranscriptFromLine(path, startLine)
@@ -409,7 +409,7 @@ func extractAgentIDFromText(text string) string {
 // CalculateTotalTokenUsage calculates token usage for a turn, including subagents.
 // It parses the main transcript from startLine, extracts spawned agent IDs,
 // and calculates their token usage from transcripts in subagentsDir.
-func CalculateTotalTokenUsage(transcriptPath string, startLine int, subagentsDir string) (*checkpoint.TokenUsage, error) {
+func CalculateTotalTokenUsage(transcriptPath string, startLine int, subagentsDir string) (*agent.TokenUsage, error) {
 	// Calculate main transcript usage
 	mainUsage, err := CalculateTokenUsageFromFile(transcriptPath, startLine)
 	if err != nil {
@@ -426,7 +426,7 @@ func CalculateTotalTokenUsage(transcriptPath string, startLine int, subagentsDir
 
 	// Calculate subagent token usage
 	if len(agentIDs) > 0 {
-		subagentUsage := &checkpoint.TokenUsage{}
+		subagentUsage := &agent.TokenUsage{}
 		for agentID := range agentIDs {
 			agentPath := filepath.Join(subagentsDir, fmt.Sprintf("agent-%s.jsonl", agentID))
 			agentUsage, err := CalculateTokenUsageFromFile(agentPath, 0)

cmd/entire/cli/agent/geminicli/gemini.go

@@ -32,11 +32,16 @@ func NewGeminiCLIAgent() agent.Agent {
 	return &GeminiCLIAgent{}
 }
 
-// Name returns the agent identifier.
-func (g *GeminiCLIAgent) Name() string {
+// Name returns the agent registry key.
+func (g *GeminiCLIAgent) Name() agent.AgentName {
 	return agent.AgentNameGemini
 }
 
+// Type returns the agent type identifier.
+func (g *GeminiCLIAgent) Type() agent.AgentType {
+	return agent.AgentTypeGemini
+}
+
 // Description returns a human-readable description.
 func (g *GeminiCLIAgent) Description() string {
 	return "Gemini CLI - Google's AI coding assistant"

cmd/entire/cli/agent/geminicli/transcript.go

@@ -4,6 +4,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+
+	"entire.io/cli/cmd/entire/cli/agent"
 )
 
 // Transcript parsing types - Gemini CLI uses JSON format for session storage
@@ -168,32 +170,20 @@ func ExtractLastAssistantMessageFromTranscript(transcript *GeminiTranscript) str
 	return ""
 }
 
-// TokenUsage represents aggregated token usage for a checkpoint
-type TokenUsage struct {
-	// InputTokens is the number of input tokens (fresh, not from cache)
-	InputTokens int `json:"input_tokens"`
-	// OutputTokens is the number of output tokens generated
-	OutputTokens int `json:"output_tokens"`
-	// CacheReadTokens is the number of tokens read from cache
-	CacheReadTokens int `json:"cache_read_tokens"`
-	// APICallCount is the number of API calls made
-	APICallCount int `json:"api_call_count"`
-}
-
 // CalculateTokenUsage calculates token usage from a Gemini transcript.
 // This is specific to Gemini's API format where each message may have a tokens object
 // with input, output, cached, thoughts, tool, and total counts.
 // Only processes messages from startMessageIndex onwards (0-indexed).
-func CalculateTokenUsage(data []byte, startMessageIndex int) *TokenUsage {
+func CalculateTokenUsage(data []byte, startMessageIndex int) *agent.TokenUsage {
 	var transcript struct {
 		Messages []geminiMessageWithTokens `json:"messages"`
 	}
 
 	if err := json.Unmarshal(data, &transcript); err != nil {
-		return &TokenUsage{}
+		return &agent.TokenUsage{}
 	}
 
-	usage := &TokenUsage{}
+	usage := &agent.TokenUsage{}
 
 	for i, msg := range transcript.Messages {
 		// Skip messages before startMessageIndex
@@ -221,15 +211,15 @@ func CalculateTokenUsage(data []byte, startMessageIndex int) *TokenUsage {
 
 // CalculateTokenUsageFromFile calculates token usage from a Gemini transcript file.
 // If startMessageIndex > 0, only considers messages from that index onwards.
-func CalculateTokenUsageFromFile(path string, startMessageIndex int) (*TokenUsage, error) {
+func CalculateTokenUsageFromFile(path string, startMessageIndex int) (*agent.TokenUsage, error) {
 	if path == "" {
-		return &TokenUsage{}, nil
+		return &agent.TokenUsage{}, nil
 	}
 
 	data, err := os.ReadFile(path) //nolint:gosec // Reading from controlled transcript path
 	if err != nil {
 		if os.IsNotExist(err) {
-			return &TokenUsage{}, nil
+			return &agent.TokenUsage{}, nil
 		}
 		return nil, fmt.Errorf("failed to read transcript: %w", err)
 	}

cmd/entire/cli/agent/registry.go

@@ -2,21 +2,21 @@ package agent
 
 import (
 	"fmt"
-	"sort"
+	"slices"
 	"sync"
 )
 
 var (
 	registryMu sync.RWMutex
-	registry   = make(map[string]Factory)
+	registry   = make(map[AgentName]Factory)
 )
 
 // Factory creates a new agent instance
 type Factory func() Agent
 
 // Register adds an agent factory to the registry.
 // Called from init() in each agent implementation.
-func Register(name string, factory Factory) {
+func Register(name AgentName, factory Factory) {
 	registryMu.Lock()
 	defer registryMu.Unlock()
 	registry[name] = factory
@@ -25,7 +25,7 @@ func Register(name string, factory Factory) {
 // Get retrieves an agent by name.
 //
 //nolint:ireturn // Factory pattern requires returning the interface
-func Get(name string) (Agent, error) {
+func Get(name AgentName) (Agent, error) {
 	registryMu.RLock()
 	defer registryMu.RUnlock()
 
@@ -37,15 +37,15 @@ func Get(name string) (Agent, error) {
 }
 
 // List returns all registered agent names in sorted order.
-func List() []string {
+func List() []AgentName {
 	registryMu.RLock()
 	defer registryMu.RUnlock()
 
-	names := make([]string, 0, len(registry))
+	names := make([]AgentName, 0, len(registry))
 	for name := range registry {
 		names = append(names, name)
 	}
-	sort.Strings(names)
+	slices.Sort(names)
 	return names
 }
 
@@ -66,40 +66,54 @@ func Detect() (Agent, error) {
 	return nil, fmt.Errorf("no agent detected (available: %v)", List())
 }
 
-// Agent name constants
+// AgentName is the registry key type for agents (e.g., "claude-code", "gemini").
+//
+//nolint:revive // stuttering is intentional - distinguishes from AgentType when both are used
+type AgentName string
+
+// AgentType is the display name type stored in metadata/trailers (e.g., "Claude Code", "Gemini CLI").
+//
+//nolint:revive // stuttering is intentional - distinguishes from AgentName when both are used
+type AgentType string
+
+// Agent name constants (registry keys)
 const (
-	AgentNameClaudeCode = "claude-code"
-	AgentNameCursor     = "cursor"
-	AgentNameWindsurf   = "windsurf"
-	AgentNameAider      = "aider"
-	AgentNameGemini     = "gemini"
+	AgentNameClaudeCode AgentName = "claude-code"
+	AgentNameGemini     AgentName = "gemini"
 )
 
-// DefaultAgentName is the default when none specified
-const DefaultAgentName = AgentNameClaudeCode
-
-// AgentTypeToRegistryName maps human-readable agent type names (as stored in session state)
-// to their registry names. Used to look up the correct agent when showing resume commands.
-var AgentTypeToRegistryName = map[string]string{
-	"Claude Code": AgentNameClaudeCode,
-	"Gemini CLI":  AgentNameGemini,
-	"Cursor":      AgentNameCursor,
-	"Windsurf":    AgentNameWindsurf,
-	"Aider":       AgentNameAider,
-}
+// Agent type constants (type identifiers stored in metadata/trailers)
+const (
+	AgentTypeClaudeCode AgentType = "Claude Code"
+	AgentTypeGemini     AgentType = "Gemini CLI"
+	AgentTypeUnknown    AgentType = "Agent" // Fallback for backwards compatibility
+)
+
+// DefaultAgentName is the registry key for the default agent.
+const DefaultAgentName AgentName = AgentNameClaudeCode
 
-// GetByAgentType retrieves an agent by its type name.
-// Accepts either human-readable names (e.g., "Claude Code", "Gemini CLI") or
-// registry names (e.g., "claude-code", "gemini").
+// GetByAgentType retrieves an agent by its type identifier.
 //
+// Note: This uses a linear search that instantiates agents until a match is found.
+// This is acceptable because:
+//   - Agent count is small (~2-20 agents)
+//   - Agent factories are lightweight (empty struct allocation)
+//   - Called infrequently (commit hooks, rewind, debug commands - not hot paths)
+//   - Cost is ~400ns worst case vs milliseconds for I/O operations
+//
+// Only optimize if agent count exceeds 100 or profiling shows this as a bottleneck.
+func GetByAgentType(agentType AgentType) (Agent, error) {
+	registryMu.RLock()
+	defer registryMu.RUnlock()
 
-func GetByAgentType(agentType string) (Agent, error) {
-	// Try human-readable name first
-	if registryName, ok := AgentTypeToRegistryName[agentType]; ok {
-		return Get(registryName)
+	for _, factory := range registry {
+		ag := factory()
+		if ag.Type() == agentType {
+			return ag, nil
+		}
 	}
-	// Fall back to treating it as a registry name
-	return Get(agentType)
+
+	return nil, fmt.Errorf("unknown agent type: %s", agentType)
 }
 
 // Default returns the default agent.