implemented ai usage tracking

app/finders/ai_usage_daily_aggregates_finder.rb

@@ -0,0 +1,39 @@
+# frozen_string_literal: true
+
+# Reads daily AI generation usage aggregates for a single owner (flow/project/namespace, or
+# the whole application) and buckets them into day/week/month periods. Callers are expected
+# to have already validated the aggregation/date range (see
+# Types::Concerns::HasAiUsageField).
+class AiUsageDailyAggregatesFinder < ApplicationFinder
+  def execute
+    params[:relation]
+      .where(date: params[:after_date]..params[:before_date])
+      .group(Arel.sql("date_trunc('#{params[:aggregation]}', date)"))
+      .order(Arel.sql("date_trunc('#{params[:aggregation]}', date)"))
+      .pluck(
+        Arel.sql("date_trunc('#{params[:aggregation]}', date)::date"),
+        Arel.sql('SUM(generation_count)'),
+        Arel.sql('SUM(total_usage)')
+      )
+      .map { |row| build_bucket(*row) }
+  end
+
+  private
+
+  def build_bucket(period_start, generation_count, total_usage)
+    Usage::Bucket.new(
+      period_start: period_start,
+      period_end: period_end_for(period_start),
+      usage: generation_count,
+      value: total_usage
+    )
+  end
+
+  def period_end_for(period_start)
+    case params[:aggregation]
+    when 'day' then period_start
+    when 'week' then period_start + 6.days
+    when 'month' then period_start.end_of_month
+    end
+  end
+end

app/finders/runtime_usage_daily_aggregates_finder.rb

@@ -21,7 +21,7 @@ def execute
   private
 
   def build_bucket(period_start, execution_count, total_execution_time_us)
-    RuntimeUsage::Bucket.new(
+    Usage::Bucket.new(
       period_start: period_start,
       period_end: period_end_for(period_start),
       usage: execution_count,

app/graphql/types/application_type.rb

@@ -3,6 +3,7 @@
 module Types
   class ApplicationType < Types::BaseObject
     include Types::Concerns::HasRuntimeUsageField
+    include Types::Concerns::HasAiUsageField
 
     description 'Represents the application instance'
 
@@ -52,6 +53,14 @@ class ApplicationType < Types::BaseObject
                         },
                         null: true
 
+    ai_usage_field description: 'Instance-wide AI generation usage, bucketed by day, week or month. ' \
+                                'Only visible to admins.',
+                   relation: ->(_object) { AiUsageDailyAggregate.all },
+                   authorized: lambda { |_object|
+                     Ability.allowed?(current_authentication, :read_application_ai_usage, :global)
+                   },
+                   null: true
+
     def metadata
       {}
     end

app/graphql/types/concerns/has_ai_usage_field.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Types
+  module Concerns
+    module HasAiUsageField
+      extend ActiveSupport::Concern
+
+      class_methods do
+        # relation: proc taking the resolved object and returning its ai_usage_daily_
+        # aggregates scope. Defaults to the association of the same name, which every owner
+        # type has.
+        # authorized: proc taking the resolved object and returning whether usage may be
+        # read; runs in the resolver instance's context (via instance_exec), so it can call
+        # instance methods like current_authentication. Defaults to always-authorized, since
+        # every current ai_usage_field caller already gates the whole type with `authorize`.
+        def ai_usage_field(description:, relation: ->(object) { object.ai_usage_daily_aggregates },
+                           authorized: ->(_object) { true }, null: false)
+          field :ai_usage, [Types::UsageBucketType], null: null, description: description do
+            argument :aggregation, Types::UsageAggregationEnum, required: false, default_value: 'day',
+                                                                description: 'Granularity to bucket usage into'
+            argument :after_date, GraphQL::Types::ISO8601Date, required: true,
+                                                               description: 'Start of the usage range (inclusive)'
+            argument :before_date, GraphQL::Types::ISO8601Date, required: true,
+                                                                description: 'End of the usage range (inclusive)'
+          end
+
+          define_method(:ai_usage) do |aggregation:, after_date:, before_date:|
+            next nil unless instance_exec(object, &authorized)
+
+            Types::Concerns::ValidatesUsageDateRange.validate_range!(aggregation: aggregation,
+                                                                     after_date: after_date, before_date: before_date)
+
+            AiUsageDailyAggregatesFinder.new(
+              relation: relation.call(object),
+              aggregation: aggregation,
+              after_date: after_date,
+              before_date: before_date
+            ).execute
+          end
+        end
+      end
+    end
+  end
+end

app/graphql/types/concerns/has_runtime_usage_field.rb

@@ -5,15 +5,6 @@ module Concerns
     module HasRuntimeUsageField
       extend ActiveSupport::Concern
 
-      # The allowed after_date..before_date span, bounded per aggregation level, per the
-      # product decision that days can be queried 7-31 days back, weeks 2-26 weeks back, and
-      # months 2-24 months back.
-      ALLOWED_SPAN = {
-        'day' => (7..31),
-        'week' => (2..26),
-        'month' => (2..24),
-      }.freeze
-
       class_methods do
         # relation: proc taking the resolved object and returning its runtime_usage_daily_
         # aggregates scope. Defaults to the association of the same name, which every owner
@@ -37,8 +28,8 @@ def runtime_usage_field(description:, relation: ->(object) { object.runtime_usag
           define_method(:runtime_usage) do |aggregation:, after_date:, before_date:|
             next nil unless instance_exec(object, &authorized)
 
-            HasRuntimeUsageField.validate_range!(aggregation: aggregation, after_date: after_date,
-                                                 before_date: before_date)
+            Types::Concerns::ValidatesUsageDateRange.validate_range!(aggregation: aggregation,
+                                                                     after_date: after_date, before_date: before_date)
 
             RuntimeUsageDailyAggregatesFinder.new(
               relation: relation.call(object),
@@ -49,37 +40,6 @@ def runtime_usage_field(description:, relation: ->(object) { object.runtime_usag
           end
         end
       end
-
-      def self.validate_range!(aggregation:, after_date:, before_date:)
-        raise GraphQL::ExecutionError, 'after_date must not be later than before_date' if after_date > before_date
-
-        allowed_span = ALLOWED_SPAN.fetch(aggregation)
-        span = span_for(aggregation, after_date, before_date)
-        return if allowed_span.cover?(span)
-
-        raise GraphQL::ExecutionError,
-              "Date range for #{aggregation} aggregation must span between " \
-              "#{allowed_span.min} and #{allowed_span.max} #{aggregation}s, got #{span}"
-      end
-
-      def self.span_for(aggregation, after_date, before_date)
-        case aggregation
-        when 'day' then (before_date - after_date).to_i + 1
-        when 'week' then weeks_between(after_date, before_date) + 1
-        when 'month' then months_between(after_date, before_date) + 1
-        end
-      end
-
-      # Counts distinct ISO weeks (Monday-start), matching date_trunc('week', ...) bucketing.
-      def self.weeks_between(from, to)
-        (to.beginning_of_week - from.beginning_of_week).to_i / 7
-      end
-
-      def self.months_between(from, to)
-        ((to.year * 12) + to.month) - ((from.year * 12) + from.month)
-      end
-
-      private_class_method :span_for, :weeks_between, :months_between
     end
   end
 end

app/graphql/types/concerns/validates_usage_date_range.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Types
+  module Concerns
+    # Shared after_date..before_date span validation for usage fields (see
+    # HasRuntimeUsageField, HasAiUsageField), per the product decision that days can be
+    # queried 7-31 days back, weeks 2-26 weeks back, and months 2-24 months back.
+    module ValidatesUsageDateRange
+      ALLOWED_SPAN = {
+        'day' => (7..31),
+        'week' => (2..26),
+        'month' => (2..24),
+      }.freeze
+
+      def self.validate_range!(aggregation:, after_date:, before_date:)
+        raise GraphQL::ExecutionError, 'after_date must not be later than before_date' if after_date > before_date
+
+        allowed_span = ALLOWED_SPAN.fetch(aggregation)
+        span = span_for(aggregation, after_date, before_date)
+        return if allowed_span.cover?(span)
+
+        raise GraphQL::ExecutionError,
+              "Date range for #{aggregation} aggregation must span between " \
+              "#{allowed_span.min} and #{allowed_span.max} #{aggregation}s, got #{span}"
+      end
+
+      def self.span_for(aggregation, after_date, before_date)
+        case aggregation
+        when 'day' then (before_date - after_date).to_i + 1
+        when 'week' then weeks_between(after_date, before_date) + 1
+        when 'month' then months_between(after_date, before_date) + 1
+        end
+      end
+
+      # Counts distinct ISO weeks (Monday-start), matching date_trunc('week', ...) bucketing.
+      def self.weeks_between(from, to)
+        (to.beginning_of_week - from.beginning_of_week).to_i / 7
+      end
+
+      def self.months_between(from, to)
+        ((to.year * 12) + to.month) - ((from.year * 12) + from.month)
+      end
+
+      private_class_method :span_for, :weeks_between, :months_between
+    end
+  end
+end

app/graphql/types/flow_type.rb

@@ -3,6 +3,7 @@
 module Types
   class FlowType < Types::BaseObject
     include Types::Concerns::HasRuntimeUsageField
+    include Types::Concerns::HasAiUsageField
 
     description 'Represents a flow'
 
@@ -68,6 +69,7 @@ class FlowType < Types::BaseObject
     ]
 
     runtime_usage_field description: 'Execution usage of this flow, bucketed by day, week or month'
+    ai_usage_field description: 'AI generation usage of this flow, bucketed by day, week or month'
 
     id_field Flow
     timestamps

app/graphql/types/namespace_project_type.rb

@@ -3,6 +3,7 @@
 module Types
   class NamespaceProjectType < Types::BaseObject
     include Types::Concerns::HasRuntimeUsageField
+    include Types::Concerns::HasAiUsageField
 
     description 'Represents a namespace project'
 
@@ -42,6 +43,7 @@ class NamespaceProjectType < Types::BaseObject
     ]
 
     runtime_usage_field description: 'Execution usage of this project, bucketed by day, week or month'
+    ai_usage_field description: 'AI generation usage of this project, bucketed by day, week or month'
 
     id_field NamespaceProject
     timestamps

app/graphql/types/namespace_type.rb

@@ -3,6 +3,7 @@
 module Types
   class NamespaceType < Types::BaseObject
     include Types::Concerns::HasRuntimeUsageField
+    include Types::Concerns::HasAiUsageField
 
     description 'Represents a Namespace'
 
@@ -36,6 +37,7 @@ class NamespaceType < Types::BaseObject
     ]
 
     runtime_usage_field description: 'Execution usage of this namespace, bucketed by day, week or month'
+    ai_usage_field description: 'AI generation usage of this namespace, bucketed by day, week or month'
 
     id_field Namespace
     timestamps

app/graphql/types/usage_aggregation_enum.rb

@@ -2,7 +2,7 @@
 
 module Types
   class UsageAggregationEnum < Types::BaseEnum
-    description 'Granularity to bucket execution usage into.'
+    description 'Granularity to bucket usage into.'
 
     value 'DAY', 'Bucket usage per day.', value: 'day'
     value 'WEEK', 'Bucket usage per ISO week.', value: 'week'

app/graphql/types/usage_bucket_type.rb

@@ -2,7 +2,7 @@
 
 module Types
   class UsageBucketType < Types::BaseObject
-    description 'Aggregated execution usage for a single day, week or month bucket'
+    description 'Aggregated usage for a single day, week or month bucket'
 
     # rubocop:disable GraphQL/ExtractType -- period_start/period_end are the bucket's own
     # range, not a nested concept worth a separate type
@@ -11,9 +11,10 @@ class UsageBucketType < Types::BaseObject
     field :period_start, GraphQL::Types::ISO8601Date, null: false,
                                                       description: 'Start date of this usage bucket (inclusive)'
     field :usage, Types::BigIntType, null: false,
-                                     description: 'Number of executions in this bucket'
+                                     description: 'Number of events (e.g. executions, ai generations) in this bucket'
     field :value, Float, null: false,
-                         description: 'Total execution time in this bucket, in seconds'
+                         description: 'Aggregated value for this bucket (e.g. execution time in ' \
+                                      'microseconds, ai usage in tokens)'
     # rubocop:enable GraphQL/ExtractType
   end
 end

app/models/ai_usage_daily_aggregate.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+# One row per project per flow per day. namespace_id is denormalized from the project so
+# namespace/application-level usage can be read by filtering (or not filtering, for
+# application-wide) this single table instead of maintaining separate per-level tables.
+class AiUsageDailyAggregate < ApplicationRecord
+  include Code0::ZeroTrack::Database::Partitioning::PartitionedTable
+  include TracksAiUsage
+
+  # Prompt-based generation of a brand-new flow has no flow yet, and a partitioned table's
+  # primary key columns can't contain NULL - rows without a real flow share this sentinel.
+  NO_FLOW = 0
+
+  partition_by :date, strategy: :monthly, retain_for: 25.months
+
+  self.table_name = 'p_ai_usage_daily_aggregates'
+  self.primary_key = %i[project_id flow_id date]
+
+  # optional: true because there's no DB-level foreign key (see the migration) - the
+  # flow/project/namespace may have been deleted while this usage row, tied to the license,
+  # must persist.
+  belongs_to :flow, inverse_of: :ai_usage_daily_aggregates, optional: true
+  belongs_to :project, class_name: 'NamespaceProject', inverse_of: :ai_usage_daily_aggregates, optional: true
+  belongs_to :namespace, inverse_of: :ai_usage_daily_aggregates, optional: true
+
+  validates :date, presence: true, uniqueness: { scope: %i[project_id flow_id] }
+end

app/models/concerns/tracks_ai_usage.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+# Atomically accumulates per-day AI generation counters via upsert, so concurrent
+# generations landing on the same owner/day never lose a counter increment to a
+# read-modify-write race. Mirrors TracksUsage's upsert shape for the AI usage metric
+# returned by Velorum (a single usage number per generation call, rather than a
+# count/duration pair).
+module TracksAiUsage
+  extend ActiveSupport::Concern
+
+  class_methods do
+    # unique_by: the conflict target's columns. Defaults to owner.keys + [:date], but callers
+    # whose owner hash also carries denormalized, non-unique columns (e.g. project_id/
+    # namespace_id alongside the real flow_id key) must pass it explicitly.
+    def record_generation!(usage:, date: Time.zone.today, unique_by: nil, **owner)
+      now = Time.current
+
+      upsert_all( # rubocop:disable Rails/SkipsModelValidations -- atomic accumulate-on-conflict upsert is the point
+        [owner.merge(date: date, generation_count: 1, total_usage: usage, created_at: now, updated_at: now)],
+        unique_by: unique_by || (owner.keys + [:date]),
+        on_duplicate: Arel.sql(<<~SQL.squish)
+          generation_count = #{table_name}.generation_count + 1,
+          total_usage = #{table_name}.total_usage + excluded.total_usage,
+          updated_at = excluded.updated_at
+        SQL
+      )
+    end
+  end
+end

app/models/flow.rb

@@ -22,6 +22,7 @@ class Flow < ApplicationRecord
   has_many :node_functions, class_name: 'NodeFunction', inverse_of: :flow
   has_many :execution_results, inverse_of: :flow
   has_many :runtime_usage_daily_aggregates, class_name: 'RuntimeUsageDailyAggregate', inverse_of: :flow
+  has_many :ai_usage_daily_aggregates, class_name: 'AiUsageDailyAggregate', inverse_of: :flow
 
   has_many :flow_data_type_links, inverse_of: :flow
   has_many :referenced_data_types, through: :flow_data_type_links, source: :referenced_data_type

app/models/namespace.rb

@@ -15,6 +15,7 @@ class Namespace < ApplicationRecord
   has_many :runtimes, inverse_of: :namespace
 
   has_many :runtime_usage_daily_aggregates, class_name: 'RuntimeUsageDailyAggregate', inverse_of: :namespace
+  has_many :ai_usage_daily_aggregates, class_name: 'AiUsageDailyAggregate', inverse_of: :namespace
 
   def organization_type?
     parent_type == Organization.name

app/models/namespace_project.rb

@@ -14,6 +14,7 @@ class NamespaceProject < ApplicationRecord
                             source: :role
   has_many :flows, class_name: 'Flow', inverse_of: :project
   has_many :runtime_usage_daily_aggregates, class_name: 'RuntimeUsageDailyAggregate', inverse_of: :project
+  has_many :ai_usage_daily_aggregates, class_name: 'AiUsageDailyAggregate', inverse_of: :project
 
   validates :slug, presence: true,
                    length: { minimum: 3, maximum: 50 },

app/policies/global_policy.rb

@@ -28,6 +28,7 @@ class GlobalPolicy < BasePolicy
     enable :list_users
     enable :create_user
     enable :read_application_usage
+    enable :read_application_ai_usage
   end
 end