diff --git a/cmd/daemon/main.go b/cmd/daemon/main.go index a80ca3f..6b07dca 100644 --- a/cmd/daemon/main.go +++ b/cmd/daemon/main.go @@ -95,6 +95,17 @@ func main() { go daemon.SocketTopicProccessor(msg) + // Start CCUsage service if enabled + if cfg.CCUsage != nil && cfg.CCUsage.Enabled != nil && *cfg.CCUsage.Enabled { + ccUsageService := model.NewCCUsageService(cfg) + if err := ccUsageService.Start(ctx); err != nil { + slog.Error("Failed to start CCUsage service", slog.Any("err", err)) + } else { + slog.Info("CCUsage service started") + defer ccUsageService.Stop() + } + } + // Create processor instance processor := daemon.NewSocketHandler(daemonConfig, pubsub) diff --git a/model/api.base.go b/model/api.base.go index cabc9fb..b376d3c 100644 --- a/model/api.base.go +++ b/model/api.base.go @@ -28,7 +28,8 @@ type HTTPRequestOptions[T any, R any] struct { Timeout time.Duration // Optional, defaults to 10 seconds } -// SendHTTPRequest is a generic HTTP request function that sends data and unmarshals the response +// SendHTTPRequest is a legacy HTTP request function that uses msgpack encoding +// Deprecated: Use SendHTTPRequestJSON for new implementations func SendHTTPRequest[T any, R any](opts HTTPRequestOptions[T, R]) error { ctx, span := modelTracer.Start(opts.Context, "http.send") defer span.End() @@ -118,3 +119,162 @@ func SendHTTPRequest[T any, R any](opts HTTPRequestOptions[T, R]) error { return nil } + +// SendHTTPRequestJSON is a generic HTTP request function that sends JSON data and unmarshals the response +func SendHTTPRequestJSON[T any, R any](opts HTTPRequestOptions[T, R]) error { + ctx, span := modelTracer.Start(opts.Context, "http.send.json") + defer span.End() + + jsonData, err := json.Marshal(opts.Payload) + if err != nil { + logrus.Errorln(err) + return err + } + + timeout := time.Second * 10 + if opts.Timeout > 0 { + timeout = opts.Timeout + } + + client := &http.Client{ + Timeout: timeout, + Transport: otelhttp.NewTransport(http.DefaultTransport), + } + + req, err := http.NewRequestWithContext(ctx, opts.Method, opts.Endpoint.APIEndpoint+opts.Path, bytes.NewBuffer(jsonData)) + if err != nil { + logrus.Errorln(err) + return err + } + + contentType := "application/json" + if opts.ContentType != "" { + contentType = opts.ContentType + } + + req.Header.Set("Content-Type", contentType) + req.Header.Set("User-Agent", fmt.Sprintf("shelltimeCLI@%s", commitID)) + req.Header.Set("Authorization", "CLI "+opts.Endpoint.Token) + + logrus.Traceln("http: ", req.URL.String()) + + resp, err := client.Do(req) + if err != nil { + logrus.Errorln(err) + return err + } + defer resp.Body.Close() + + logrus.Traceln("http: ", resp.Status) + + if resp.StatusCode == http.StatusNoContent { + return nil + } + + buf, err := io.ReadAll(resp.Body) + if err != nil { + logrus.Errorln(err) + return err + } + + if resp.StatusCode != http.StatusOK { + var msg errorResponse + err = json.Unmarshal(buf, &msg) + if err != nil { + logrus.Errorln("Failed to parse error response:", err) + return fmt.Errorf("HTTP error: %d", resp.StatusCode) + } + logrus.Errorln("Error response:", msg.ErrorMessage) + return errors.New(msg.ErrorMessage) + } + + // Only try to unmarshal if we have a response struct + if opts.Response != nil { + err = json.Unmarshal(buf, opts.Response) + if err != nil { + logrus.Errorln("Failed to unmarshal JSON response:", err) + return err + } + } + + return nil +} + +// GraphQLResponse is a generic wrapper for GraphQL responses +type GraphQLResponse[T any] struct { + Data T `json:"data"` + Errors []GraphQLError `json:"errors,omitempty"` +} + +// GraphQLError represents a GraphQL error +type GraphQLError struct { + Message string `json:"message"` + Extensions map[string]interface{} `json:"extensions,omitempty"` + Path []interface{} `json:"path,omitempty"` +} + +// GraphQLRequestOptions contains options for GraphQL requests +type GraphQLRequestOptions[R any] struct { + Context context.Context + Endpoint Endpoint + Query string + Variables map[string]interface{} + Response *R + Timeout time.Duration // Optional, defaults to 30 seconds +} + +// SendGraphQLRequest sends a GraphQL request and unmarshals the response +func SendGraphQLRequest[R any](opts GraphQLRequestOptions[R]) error { + ctx, span := modelTracer.Start(opts.Context, "graphql.send") + defer span.End() + + // Build GraphQL payload + payload := map[string]interface{}{ + "query": opts.Query, + } + if opts.Variables != nil { + payload["variables"] = opts.Variables + } + + // Build GraphQL endpoint path + graphQLPath := "/api/v2/graphql" + + // Set default timeout + timeout := time.Second * 30 + if opts.Timeout > 0 { + timeout = opts.Timeout + } + + // Use the new JSON HTTP request function + err := SendHTTPRequestJSON(HTTPRequestOptions[map[string]interface{}, R]{ + Context: ctx, + Endpoint: opts.Endpoint, + Method: http.MethodPost, + Path: graphQLPath, + Payload: payload, + Response: opts.Response, + Timeout: timeout, + }) + + if err != nil { + // The error is already formatted by SendHTTPRequestJSON + return err + } + + // Check for GraphQL errors in the response if we have a response + if opts.Response != nil { + // Marshal response back to check for errors + respBytes, err := json.Marshal(opts.Response) + if err == nil { + var errorCheck struct { + Errors []GraphQLError `json:"errors,omitempty"` + } + if err := json.Unmarshal(respBytes, &errorCheck); err == nil && len(errorCheck.Errors) > 0 { + // Return the first error message if there are GraphQL errors + return fmt.Errorf("GraphQL error: %s", errorCheck.Errors[0].Message) + } + } + } + + return nil +} diff --git a/model/ccusage_service.go b/model/ccusage_service.go new file mode 100644 index 0000000..f52b2e8 --- /dev/null +++ b/model/ccusage_service.go @@ -0,0 +1,409 @@ +package model + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "os" + "os/exec" + "os/user" + "time" + + "github.com/sirupsen/logrus" +) + +// CCUsageData represents the usage data collected from ccusage command +type CCUsageData struct { + Timestamp string `json:"timestamp" msgpack:"timestamp"` + Hostname string `json:"hostname" msgpack:"hostname"` + Username string `json:"username" msgpack:"username"` + OS string `json:"os" msgpack:"os"` + OSVersion string `json:"osVersion" msgpack:"osVersion"` + Data CCUsageProjectDailyOutput `json:"data" msgpack:"data"` +} + +// CCUsageService defines the interface for CC usage collection +type CCUsageService interface { + Start(ctx context.Context) error + Stop() + CollectCCUsage(ctx context.Context) error +} + +// ccUsageService implements the CCUsageService interface +type ccUsageService struct { + config ShellTimeConfig + ticker *time.Ticker + stopChan chan struct{} +} + +// NewCCUsageService creates a new CCUsage service +func NewCCUsageService(config ShellTimeConfig) CCUsageService { + return &ccUsageService{ + config: config, + stopChan: make(chan struct{}), + } +} + +// Start begins the periodic usage collection +func (s *ccUsageService) Start(ctx context.Context) error { + // Check if CCUsage is enabled + if s.config.CCUsage == nil || s.config.CCUsage.Enabled == nil || !*s.config.CCUsage.Enabled { + logrus.Info("CCUsage collection is disabled") + return nil + } + + logrus.Info("Starting CCUsage collection service") + + // Create a ticker for hourly collection + s.ticker = time.NewTicker(1 * time.Hour) + + // Run initial collection + if err := s.CollectCCUsage(ctx); err != nil { + logrus.Warnf("Initial CCUsage collection failed: %v", err) + } + + // Start the collection loop + go func() { + for { + select { + case <-s.ticker.C: + if err := s.CollectCCUsage(ctx); err != nil { + logrus.Warnf("CCUsage collection failed: %v", err) + } + case <-s.stopChan: + logrus.Info("Stopping CCUsage collection service") + return + case <-ctx.Done(): + logrus.Info("Context cancelled, stopping CCUsage collection service") + return + } + } + }() + + return nil +} + +// Stop halts the usage collection +func (s *ccUsageService) Stop() { + if s.ticker != nil { + s.ticker.Stop() + } + close(s.stopChan) +} + +// CollectCCUsage collects and sends usage data to the server +func (s *ccUsageService) CollectCCUsage(ctx context.Context) error { + ctx, span := modelTracer.Start(ctx, "ccusage.collect") + defer span.End() + + logrus.Debug("Collecting CCUsage data") + + since := time.Time{} + + // Get the last sync timestamp from server if we have credentials + if s.config.Token != "" && s.config.APIEndpoint != "" { + endpoint := Endpoint{ + Token: s.config.Token, + APIEndpoint: s.config.APIEndpoint, + } + + // Try to get last sync timestamp, but don't fail if it doesn't work + lastSync, err := s.getLastSyncTimestamp(ctx, endpoint) + if err != nil { + logrus.Warnf("Failed to get last sync timestamp: %v", err) + } + since = lastSync + logrus.Debugf("Got last sync timestamp: %v\n", since) + } + + // Collect data from ccusage command + data, err := s.collectData(ctx, since) + if err != nil { + return fmt.Errorf("failed to collect ccusage data: %w", err) + } + + // Send to server + if s.config.Token != "" && s.config.APIEndpoint != "" { + endpoint := Endpoint{ + Token: s.config.Token, + APIEndpoint: s.config.APIEndpoint, + } + + err = s.sendData(ctx, endpoint, data) + if err != nil { + return fmt.Errorf("failed to send usage data: %w", err) + } + } + + logrus.Debug("CCUsage data collection completed") + return nil +} + +// getLastSyncTimestamp fetches the last CCUsage sync timestamp from the server via GraphQL +func (s *ccUsageService) getLastSyncTimestamp(ctx context.Context, endpoint Endpoint) (time.Time, error) { + // Get current hostname + hostname, err := os.Hostname() + if err != nil { + logrus.Warnf("Failed to get hostname: %v", err) + hostname = "unknown" + } + + query := `query fetchUserCCUsageLastSync($hostname: String!) { + fetchUser { + id + ccusage(filter: { hostname: $hostname }) { + lastSyncAt + } + } + }` + + type fetchUserResponse struct { + FetchUser struct { + ID int `json:"id"` + CCUsage struct { + LastSyncAt string `json:"lastSyncAt"` + } `json:"ccusage"` + } `json:"fetchUser"` + } + + var result GraphQLResponse[fetchUserResponse] + + variables := map[string]interface{}{ + "hostname": hostname, + } + + logrus.Debugf("Fetching CCUsage last sync for hostname: %s", hostname) + + err = SendGraphQLRequest(GraphQLRequestOptions[GraphQLResponse[fetchUserResponse]]{ + Context: ctx, + Endpoint: endpoint, + Query: query, + Variables: variables, + Response: &result, + Timeout: time.Second * 10, + }) + + if err != nil { + logrus.Warnf("Failed to fetch CCUsage last sync: %v", err) + return time.Time{}, nil // Return nil to skip the since parameter + } + + lastSyncAtStr := result.Data.FetchUser.CCUsage.LastSyncAt + + if lastSyncAtStr == "" { + return time.Time{}, nil + } + lastSyncAt, err := time.Parse(time.RFC3339, lastSyncAtStr) + if err != nil { + logrus.Warnf("Failed to parse last sync timestamp: %v", err) + return time.Time{}, err // Return nil to skip the since parameter + } + + year2023 := time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC) + if lastSyncAt.Before(year2023) { + return time.Time{}, nil + } + + return lastSyncAt, nil +} + +// collectData collects usage data using bunx or npx ccusage command +func (s *ccUsageService) collectData(ctx context.Context, since time.Time) (*CCUsageData, error) { + // Check if bunx exists + bunxPath, bunxErr := exec.LookPath("bunx") + npxPath, npxErr := exec.LookPath("npx") + + if bunxErr != nil && npxErr != nil { + return nil, fmt.Errorf("neither bunx nor npx found in system PATH") + } + + // Build command arguments + args := []string{"ccusage", "daily", "--instances", "--json"} + + // Add since parameter if provided + if !since.IsZero() { + // Convert Unix timestamp (seconds) to ISO 8601 date string + sinceDate := since.Format("20060102") + args = append(args, "--since", sinceDate) + logrus.Debugf("Using since parameter: %s (from timestamp %v)\n", sinceDate, since) + } + + var cmd *exec.Cmd + if bunxErr == nil { + // Use bunx if available + cmd = exec.CommandContext(ctx, bunxPath, args...) + logrus.Debug("Using bunx to collect ccusage data") + } else { + // Fall back to npx + cmd = exec.CommandContext(ctx, npxPath, args...) + logrus.Debug("Using npx to collect ccusage data") + } + + // Execute the command + output, err := cmd.Output() + if err != nil { + if exitErr, ok := err.(*exec.ExitError); ok { + return nil, fmt.Errorf("ccusage command failed: %v, stderr: %s", err, string(exitErr.Stderr)) + } + return nil, fmt.Errorf("failed to execute ccusage command: %w", err) + } + + // Parse JSON output + var ccusageOutput CCUsageProjectDailyOutput + if err := json.Unmarshal(output, &ccusageOutput); err != nil { + return nil, fmt.Errorf("failed to parse ccusage output: %w", err) + } + + // Get system information for metadata + hostname, err := os.Hostname() + if err != nil { + logrus.Warnf("Failed to get hostname: %v", err) + hostname = "unknown" + } + + username := os.Getenv("USER") + if username == "" { + currentUser, err := user.Current() + if err != nil { + logrus.Warnf("Failed to get username: %v", err) + username = "unknown" + } else { + username = currentUser.Username + } + } + + sysInfo, err := GetOSAndVersion() + if err != nil { + logrus.Warnf("Failed to get OS info: %v", err) + sysInfo = &SysInfo{ + Os: "unknown", + Version: "unknown", + } + } + + data := &CCUsageData{ + Timestamp: time.Now().Format(time.RFC3339), + Hostname: hostname, + Username: username, + OS: sysInfo.Os, + OSVersion: sysInfo.Version, + Data: ccusageOutput, + } + + return data, nil +} + +// sendData sends the collected usage data to the server +func (s *ccUsageService) sendData(ctx context.Context, endpoint Endpoint, data *CCUsageData) error { + // CCUsage batch request types matching server handler + type ccUsageModelBreakdown struct { + ModelName string `json:"modelName"` + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + CacheCreationTokens int `json:"cacheCreationTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + Cost float64 `json:"cost"` + } + + type ccUsageDailyData struct { + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + CacheCreationTokens int `json:"cacheCreationTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + TotalTokens int `json:"totalTokens"` + TotalCost float64 `json:"totalCost"` + ModelsUsed []string `json:"modelsUsed"` + ModelBreakdowns []ccUsageModelBreakdown `json:"modelBreakdowns"` + } + + type ccUsageEntry struct { + Project string `json:"project"` + Date string `json:"date"` // YYYYMMDD format + Usage ccUsageDailyData `json:"usage"` + } + + type ccUsageBatchPayload struct { + Host string `json:"host"` + Entries []ccUsageEntry `json:"entries"` + } + + type ccUsageResponse struct { + Success bool `json:"success"` + SuccessCount int `json:"successCount"` + TotalCount int `json:"totalCount"` + FailedProjects []string `json:"failedProjects,omitempty"` + } + + // Transform CCUsageData to batch format + var entries []ccUsageEntry + + // Iterate through all projects in the collected data + for projectName, projectDays := range data.Data.Projects { + for _, dayData := range projectDays { + // Convert model breakdowns + modelBreakdowns := make([]ccUsageModelBreakdown, len(dayData.ModelBreakdowns)) + for i, mb := range dayData.ModelBreakdowns { + modelBreakdowns[i] = ccUsageModelBreakdown{ + ModelName: mb.ModelName, + InputTokens: mb.InputTokens, + OutputTokens: mb.OutputTokens, + CacheCreationTokens: mb.CacheCreationTokens, + CacheReadTokens: mb.CacheReadTokens, + Cost: mb.Cost, + } + } + + entry := ccUsageEntry{ + Project: projectName, + Date: dayData.Date, // Already in YYYYMMDD format from ccusage + Usage: ccUsageDailyData{ + InputTokens: dayData.InputTokens, + OutputTokens: dayData.OutputTokens, + CacheCreationTokens: dayData.CacheCreationTokens, + CacheReadTokens: dayData.CacheReadTokens, + TotalTokens: dayData.TotalTokens, + TotalCost: dayData.TotalCost, + ModelsUsed: dayData.ModelsUsed, + ModelBreakdowns: modelBreakdowns, + }, + } + entries = append(entries, entry) + } + } + + if len(entries) == 0 { + logrus.Debug("No CCUsage entries to send") + return nil + } + + payload := ccUsageBatchPayload{ + Host: data.Hostname, + Entries: entries, + } + + var resp ccUsageResponse + + err := SendHTTPRequestJSON(HTTPRequestOptions[ccUsageBatchPayload, ccUsageResponse]{ + Context: ctx, + Endpoint: endpoint, + Method: http.MethodPost, + Path: "/api/v1/ccusage/batch", + Payload: payload, + Response: &resp, + }) + + if err != nil { + return fmt.Errorf("failed to send CCUsage data: %w", err) + } + + if !resp.Success { + if len(resp.FailedProjects) > 0 { + return fmt.Errorf("server rejected CCUsage data for projects: %v", resp.FailedProjects) + } + return fmt.Errorf("server rejected CCUsage data: %d/%d entries failed", resp.TotalCount-resp.SuccessCount, resp.TotalCount) + } + + logrus.Debugf("CCUsage data sent successfully: %d/%d entries", resp.SuccessCount, resp.TotalCount) + return nil +} diff --git a/model/ccusage_service.types.go b/model/ccusage_service.types.go new file mode 100644 index 0000000..3b1f61b --- /dev/null +++ b/model/ccusage_service.types.go @@ -0,0 +1,30 @@ +package model + +type CCUsageProjectDailyOutput struct { + Projects map[string][]struct { + Date string `json:"date"` + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + CacheCreationTokens int `json:"cacheCreationTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + TotalTokens int `json:"totalTokens"` + TotalCost float64 `json:"totalCost"` + ModelsUsed []string `json:"modelsUsed"` + ModelBreakdowns []struct { + ModelName string `json:"modelName"` + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + CacheCreationTokens int `json:"cacheCreationTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + Cost float64 `json:"cost"` + } `json:"modelBreakdowns"` + } `json:"projects"` + Totals struct { + InputTokens int `json:"inputTokens"` + OutputTokens int `json:"outputTokens"` + CacheCreationTokens int `json:"cacheCreationTokens"` + CacheReadTokens int `json:"cacheReadTokens"` + TotalCost float64 `json:"totalCost"` + TotalTokens int `json:"totalTokens"` + } `json:"totals"` +} diff --git a/model/dotfile.go b/model/dotfile.go index 3f03d43..6c79e73 100644 --- a/model/dotfile.go +++ b/model/dotfile.go @@ -1,14 +1,10 @@ package model import ( - "bytes" "context" - "encoding/json" "fmt" - "io" "net/http" "os" - "strings" "time" "github.com/sirupsen/logrus" @@ -133,18 +129,19 @@ type DotfileAppResponse struct { Files []DotfileFile `json:"files"` } -type FetchDotfilesResponse struct { - Data struct { - FetchUser struct { - ID int `json:"id"` - Dotfiles struct { - TotalCount int `json:"totalCount"` - Apps []DotfileAppResponse `json:"apps"` - } `json:"dotfiles"` - } `json:"fetchUser"` - } `json:"data"` +type FetchUserDotfilesData struct { + FetchUser struct { + ID int `json:"id"` + Dotfiles struct { + TotalCount int `json:"totalCount"` + Apps []DotfileAppResponse `json:"apps"` + } `json:"dotfiles"` + } `json:"fetchUser"` } +// FetchDotfilesResponse is the complete GraphQL response for dotfiles +type FetchDotfilesResponse = GraphQLResponse[FetchUserDotfilesData] + // FetchDotfilesFromServer fetches dotfiles from the server using GraphQL func FetchDotfilesFromServer(ctx context.Context, endpoint Endpoint, filter *DotfileFilter) (*FetchDotfilesResponse, error) { query := `query fetchUserDotfiles($filter: DotfileFilter) { @@ -183,55 +180,19 @@ func FetchDotfilesFromServer(ctx context.Context, endpoint Endpoint, filter *Dot variables["filter"] = filter } - payload := map[string]interface{}{ - "query": query, - "variables": variables, - } - - jsonData, err := json.Marshal(payload) - if err != nil { - return nil, err - } - - client := &http.Client{ - Timeout: time.Second * 30, - } - - // Use web endpoint for GraphQL queries - graphQLEndpoint := endpoint.APIEndpoint - graphQLEndpoint = strings.TrimSuffix(graphQLEndpoint, "/") - if !strings.HasSuffix(graphQLEndpoint, "/api/v2/graphql") { - graphQLEndpoint += "/api/v2/graphql" - } - - req, err := http.NewRequestWithContext(ctx, http.MethodPost, graphQLEndpoint, bytes.NewBuffer(jsonData)) - if err != nil { - return nil, err - } - - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "CLI "+endpoint.Token) - req.Header.Set("User-Agent", fmt.Sprintf("shelltimeCLI@%s", commitID)) - - resp, err := client.Do(req) - if err != nil { - return nil, err - } - defer resp.Body.Close() + var result FetchDotfilesResponse + err := SendGraphQLRequest(GraphQLRequestOptions[FetchDotfilesResponse]{ + Context: ctx, + Endpoint: endpoint, + Query: query, + Variables: variables, + Response: &result, + Timeout: time.Second * 30, + }) - body, err := io.ReadAll(resp.Body) if err != nil { return nil, err } - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("GraphQL request failed with status %d: %s", resp.StatusCode, string(body)) - } - - var result FetchDotfilesResponse - if err := json.Unmarshal(body, &result); err != nil { - return nil, fmt.Errorf("failed to parse GraphQL response: %w", err) - } - return &result, nil } diff --git a/model/types.go b/model/types.go index 2730412..4a6db7a 100644 --- a/model/types.go +++ b/model/types.go @@ -16,6 +16,10 @@ type AIConfig struct { Agent AIAgentConfig `toml:"agent"` } +type CCUsage struct { + Enabled *bool `toml:"enabled"` +} + type ShellTimeConfig struct { Token string APIEndpoint string @@ -47,6 +51,9 @@ type ShellTimeConfig struct { // Exclude patterns - regular expressions to exclude commands from being saved // Commands matching any of these patterns will not be synced to the server Exclude []string `toml:"exclude"` + + // CCUsage configuration for Claude Code usage tracking + CCUsage *CCUsage `toml:"ccusage"` } var DefaultAIConfig = &AIConfig{ @@ -70,4 +77,5 @@ var DefaultConfig = ShellTimeConfig{ Encrypted: nil, AI: DefaultAIConfig, Exclude: []string{}, + CCUsage: nil, }