Merge pull request #106 from shelltime/feat/ccusage-collection-service

feat(ccusage): add CCUsage collection service for usage tracking

Author: AnnatarHe

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)
 

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
+}