diff --git a/app/finders/ai_usage_daily_aggregates_finder.rb b/app/finders/ai_usage_daily_aggregates_finder.rb new file mode 100644 index 00000000..98fe44cf --- /dev/null +++ b/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 diff --git a/app/finders/runtime_usage_daily_aggregates_finder.rb b/app/finders/runtime_usage_daily_aggregates_finder.rb index f4c3ba44..e4dfd951 100644 --- a/app/finders/runtime_usage_daily_aggregates_finder.rb +++ b/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, diff --git a/app/graphql/types/application_type.rb b/app/graphql/types/application_type.rb index d9ed048f..552fb9a4 100644 --- a/app/graphql/types/application_type.rb +++ b/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 diff --git a/app/graphql/types/concerns/has_ai_usage_field.rb b/app/graphql/types/concerns/has_ai_usage_field.rb new file mode 100644 index 00000000..4645bad8 --- /dev/null +++ b/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 diff --git a/app/graphql/types/concerns/has_runtime_usage_field.rb b/app/graphql/types/concerns/has_runtime_usage_field.rb index 1978b98a..5442cd97 100644 --- a/app/graphql/types/concerns/has_runtime_usage_field.rb +++ b/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 diff --git a/app/graphql/types/concerns/validates_usage_date_range.rb b/app/graphql/types/concerns/validates_usage_date_range.rb new file mode 100644 index 00000000..5ae34f1b --- /dev/null +++ b/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 diff --git a/app/graphql/types/flow_type.rb b/app/graphql/types/flow_type.rb index a643dc87..2e7c1f13 100644 --- a/app/graphql/types/flow_type.rb +++ b/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 diff --git a/app/graphql/types/namespace_project_type.rb b/app/graphql/types/namespace_project_type.rb index a59e6c93..8a3e5b59 100644 --- a/app/graphql/types/namespace_project_type.rb +++ b/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 diff --git a/app/graphql/types/namespace_type.rb b/app/graphql/types/namespace_type.rb index a1ac1cc8..9bf96420 100644 --- a/app/graphql/types/namespace_type.rb +++ b/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 diff --git a/app/graphql/types/usage_aggregation_enum.rb b/app/graphql/types/usage_aggregation_enum.rb index c8b72191..8d575f48 100644 --- a/app/graphql/types/usage_aggregation_enum.rb +++ b/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' diff --git a/app/graphql/types/usage_bucket_type.rb b/app/graphql/types/usage_bucket_type.rb index 34905c48..d0ca16f3 100644 --- a/app/graphql/types/usage_bucket_type.rb +++ b/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 diff --git a/app/models/ai_usage_daily_aggregate.rb b/app/models/ai_usage_daily_aggregate.rb new file mode 100644 index 00000000..92517d71 --- /dev/null +++ b/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 diff --git a/app/models/concerns/tracks_ai_usage.rb b/app/models/concerns/tracks_ai_usage.rb new file mode 100644 index 00000000..2f650eee --- /dev/null +++ b/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 diff --git a/app/models/flow.rb b/app/models/flow.rb index be1f18c2..6204efc0 100644 --- a/app/models/flow.rb +++ b/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 diff --git a/app/models/namespace.rb b/app/models/namespace.rb index ac46be84..09731302 100644 --- a/app/models/namespace.rb +++ b/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 diff --git a/app/models/namespace_project.rb b/app/models/namespace_project.rb index b6437ea8..13baa933 100644 --- a/app/models/namespace_project.rb +++ b/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 }, diff --git a/app/policies/global_policy.rb b/app/policies/global_policy.rb index e5a57225..78b4fd61 100644 --- a/app/policies/global_policy.rb +++ b/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 diff --git a/app/services/error_code.rb b/app/services/error_code.rb index 379dff33..2d81596b 100644 --- a/app/services/error_code.rb +++ b/app/services/error_code.rb @@ -111,6 +111,7 @@ def self.error_codes invalid_execution_result: { description: 'The execution result is invalid because of active model errors' }, lock_timeout: { description: 'Could not acquire a database lock in time' }, execution_usage_recording_failed: { description: 'Failed to record execution usage counters' }, + generation_usage_recording_failed: { description: 'Failed to record AI generation usage counters' }, } end # rubocop:enable Layout/LineLength diff --git a/app/services/runtime_usage/bucket.rb b/app/services/usage/bucket.rb similarity index 85% rename from app/services/runtime_usage/bucket.rb rename to app/services/usage/bucket.rb index 99e8d22b..5f77836c 100644 --- a/app/services/runtime_usage/bucket.rb +++ b/app/services/usage/bucket.rb @@ -1,5 +1,5 @@ # frozen_string_literal: true -module RuntimeUsage +module Usage Bucket = Struct.new(:period_start, :period_end, :usage, :value, keyword_init: true) end diff --git a/app/services/velorum/generate_flow_service.rb b/app/services/velorum/generate_flow_service.rb index 77364bc9..06f97fa6 100644 --- a/app/services/velorum/generate_flow_service.rb +++ b/app/services/velorum/generate_flow_service.rb @@ -45,6 +45,8 @@ def execute serialized_flow = GenerationFlowSerializer.new(response.flow, project: project).to_h + RecordGenerationUsageService.new(project: project, usage: response.usage, flow: flow).execute + ServiceResponse.success( message: 'Generated flow', payload: { diff --git a/app/services/velorum/record_generation_usage_service.rb b/app/services/velorum/record_generation_usage_service.rb new file mode 100644 index 00000000..f60b66f5 --- /dev/null +++ b/app/services/velorum/record_generation_usage_service.rb @@ -0,0 +1,35 @@ +# frozen_string_literal: true + +module Velorum + # Rolls a successful generation call up into its project's daily AI usage row. Runs after + # GenerateFlowService's own success path so a failure here never blocks returning the + # generated flow to the caller: any error is rescued and logged instead of propagating. + # namespace_id is stored denormalized on the same row so namespace/application-level usage + # can be read straight off this one table. + class RecordGenerationUsageService + include Code0::ZeroTrack::Loggable + + attr_reader :project, :usage, :flow + + def initialize(project:, usage:, flow: nil) + @project = project + @usage = usage + @flow = flow + end + + def execute + AiUsageDailyAggregate.record_generation!( + project_id: project.id, namespace_id: project.namespace_id, + flow_id: flow&.id || AiUsageDailyAggregate::NO_FLOW, + usage: usage, unique_by: %i[project_id flow_id date] + ) + + ServiceResponse.success(message: 'Generation usage recorded') + rescue StandardError => e + logger.error(message: 'Failed to record generation usage', project_id: project.id, error: e.message) + + ServiceResponse.error(message: 'Failed to record generation usage', + error_code: :generation_usage_recording_failed) + end + end +end diff --git a/config/initializers/1_zero_track.rb b/config/initializers/1_zero_track.rb index d069bed8..da9f43a4 100644 --- a/config/initializers/1_zero_track.rb +++ b/config/initializers/1_zero_track.rb @@ -18,6 +18,7 @@ partition_manager.register_model(RuntimeModuleStatusDailyUptime) partition_manager.register_model(RuntimeStatusDailyUptime) partition_manager.register_model(RuntimeUsageDailyAggregate) + partition_manager.register_model(AiUsageDailyAggregate) end config.after_initialize do diff --git a/db/migrate/20260815100000_create_ai_usage_daily_aggregates.rb b/db/migrate/20260815100000_create_ai_usage_daily_aggregates.rb new file mode 100644 index 00000000..14926dd6 --- /dev/null +++ b/db/migrate/20260815100000_create_ai_usage_daily_aggregates.rb @@ -0,0 +1,30 @@ +# frozen_string_literal: true + +class CreateAiUsageDailyAggregates < Code0::ZeroTrack::Database::Migration[1.0] + def change + # One row per project per flow per day. flow_id defaults to 0 (NO_FLOW sentinel, see + # AiUsageDailyAggregate) rather than being nullable, because a partitioned table's + # primary key columns can't contain NULL - prompt-based generation of a brand-new flow + # has no flow_id yet, so those rows share the sentinel instead of a real flow. + # namespace_id is denormalized from the project at write time so namespace/application- + # level usage can be read by filtering (or not filtering, for application-wide) this + # single table instead of writing multiple rows per generation. + # + # No foreign keys on flow_id/project_id/namespace_id: usage rows are tied to the license + # and must outlive their flow/project/namespace, so deleting those records must never + # cascade or nullify here (flow_id/project_id are also part of the primary key, so + # nullifying them would violate the PK anyway). + create_partition_by_date_table :p_ai_usage_daily_aggregates, + partition_column: :date, + primary_key: %i[project_id flow_id date] do |t| + t.bigint :flow_id, null: false, default: 0 + t.references :project, null: false, index: false + t.references :namespace, null: false + t.date :date, null: false + t.bigint :generation_count, null: false, default: 0 + t.bigint :total_usage, null: false, default: 0 + + t.timestamps_with_timezone + end + end +end diff --git a/db/schema_migrations/20260815100000 b/db/schema_migrations/20260815100000 new file mode 100644 index 00000000..3cc2809e --- /dev/null +++ b/db/schema_migrations/20260815100000 @@ -0,0 +1 @@ +8623638bff23a6501949620f0708c6403ba363665ab1e9ce2ef0f4aeb4683bd6 \ No newline at end of file diff --git a/db/structure.sql b/db/structure.sql index 3dfe2d00..8153f0cd 100644 --- a/db/structure.sql +++ b/db/structure.sql @@ -684,6 +684,18 @@ CREATE SEQUENCE organizations_id_seq ALTER SEQUENCE organizations_id_seq OWNED BY organizations.id; +CREATE TABLE p_ai_usage_daily_aggregates ( + flow_id bigint DEFAULT 0 NOT NULL, + project_id bigint NOT NULL, + namespace_id bigint NOT NULL, + date date NOT NULL, + generation_count bigint DEFAULT 0 NOT NULL, + total_usage bigint DEFAULT 0 NOT NULL, + created_at timestamp with time zone NOT NULL, + updated_at timestamp with time zone NOT NULL +) +PARTITION BY RANGE (date); + CREATE TABLE p_audit_events ( id bigint NOT NULL, author_id bigint NOT NULL, @@ -1514,6 +1526,9 @@ ALTER TABLE ONLY node_parameters ALTER TABLE ONLY organizations ADD CONSTRAINT organizations_pkey PRIMARY KEY (id); +ALTER TABLE ONLY p_ai_usage_daily_aggregates + ADD CONSTRAINT p_ai_usage_daily_aggregates_pkey PRIMARY KEY (project_id, flow_id, date); + ALTER TABLE ONLY p_audit_events ADD CONSTRAINT p_audit_events_pkey PRIMARY KEY (id, created_at); @@ -1774,6 +1789,8 @@ CREATE INDEX index_node_parameters_on_parameter_definition_id ON node_parameters CREATE UNIQUE INDEX "index_organizations_on_LOWER_name" ON organizations USING btree (lower(name)); +CREATE INDEX index_p_ai_usage_daily_aggregates_on_namespace_id ON ONLY p_ai_usage_daily_aggregates USING btree (namespace_id); + CREATE INDEX index_p_audit_events_on_author_id ON ONLY p_audit_events USING btree (author_id); CREATE INDEX index_p_execution_node_results_on_function_definition_id ON ONLY p_execution_node_results USING btree (function_definition_id); diff --git a/docs/graphql/enum/errorcodeenum.md b/docs/graphql/enum/errorcodeenum.md index 04c3cdb9..d542e610 100644 --- a/docs/graphql/enum/errorcodeenum.md +++ b/docs/graphql/enum/errorcodeenum.md @@ -24,6 +24,7 @@ Represents the available error responses | `FLOW_GENERATION_FAILED` | Flow generation failed | | `FLOW_NOT_FOUND` | The flow with the given identifier was not found | | `FLOW_TYPE_NOT_FOUND` | The flow type with the given identifier was not found | +| `GENERATION_USAGE_RECORDING_FAILED` | Failed to record AI generation usage counters | | `GENERIC_KEY_NOT_FOUND` | The given key was not found in the data type | | `IDENTITY_NOT_FOUND` | The external identity with the given identifier was not found | | `IDENTITY_VALIDATION_FAILED` | Failed to validate the external identity | diff --git a/docs/graphql/enum/usageaggregation.md b/docs/graphql/enum/usageaggregation.md index c20b4ec1..b1a48214 100644 --- a/docs/graphql/enum/usageaggregation.md +++ b/docs/graphql/enum/usageaggregation.md @@ -2,7 +2,7 @@ title: UsageAggregation --- -Granularity to bucket execution usage into. +Granularity to bucket usage into. | Value | Description | |-------|-------------| diff --git a/docs/graphql/object/application.md b/docs/graphql/object/application.md index d66945ba..79f37f5a 100644 --- a/docs/graphql/object/application.md +++ b/docs/graphql/object/application.md @@ -20,6 +20,18 @@ Represents the application instance ## Fields with arguments +### aiUsage + +Instance-wide AI generation usage, bucketed by day, week or month. Only visible to admins. + +Returns [`[UsageBucket!]`](../object/usagebucket.md). + +| Name | Type | Description | +|------|------|-------------| +| `afterDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | Start of the usage range (inclusive) | +| `aggregation` | [`UsageAggregation`](../enum/usageaggregation.md) | Granularity to bucket usage into | +| `beforeDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | End of the usage range (inclusive) | + ### identityProviderLoginUrl Login URL for a specific identity provider diff --git a/docs/graphql/object/flow.md b/docs/graphql/object/flow.md index b4a8d0ee..6a8f17ee 100644 --- a/docs/graphql/object/flow.md +++ b/docs/graphql/object/flow.md @@ -27,6 +27,18 @@ Represents a flow ## Fields with arguments +### aiUsage + +AI generation usage of this flow, bucketed by day, week or month + +Returns [`[UsageBucket!]!`](../object/usagebucket.md). + +| Name | Type | Description | +|------|------|-------------| +| `afterDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | Start of the usage range (inclusive) | +| `aggregation` | [`UsageAggregation`](../enum/usageaggregation.md) | Granularity to bucket usage into | +| `beforeDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | End of the usage range (inclusive) | + ### executionResult Find an execution result by runtime identifier diff --git a/docs/graphql/object/namespace.md b/docs/graphql/object/namespace.md index 1f7032d0..d4b64873 100644 --- a/docs/graphql/object/namespace.md +++ b/docs/graphql/object/namespace.md @@ -22,6 +22,18 @@ Represents a Namespace ## Fields with arguments +### aiUsage + +AI generation usage of this namespace, bucketed by day, week or month + +Returns [`[UsageBucket!]!`](../object/usagebucket.md). + +| Name | Type | Description | +|------|------|-------------| +| `afterDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | Start of the usage range (inclusive) | +| `aggregation` | [`UsageAggregation`](../enum/usageaggregation.md) | Granularity to bucket usage into | +| `beforeDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | End of the usage range (inclusive) | + ### project Query a project by its id diff --git a/docs/graphql/object/namespaceproject.md b/docs/graphql/object/namespaceproject.md index ea2b96c3..b89208f6 100644 --- a/docs/graphql/object/namespaceproject.md +++ b/docs/graphql/object/namespaceproject.md @@ -24,6 +24,18 @@ Represents a namespace project ## Fields with arguments +### aiUsage + +AI generation usage of this project, bucketed by day, week or month + +Returns [`[UsageBucket!]!`](../object/usagebucket.md). + +| Name | Type | Description | +|------|------|-------------| +| `afterDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | Start of the usage range (inclusive) | +| `aggregation` | [`UsageAggregation`](../enum/usageaggregation.md) | Granularity to bucket usage into | +| `beforeDate` | [`ISO8601Date!`](../scalar/iso8601date.md) | End of the usage range (inclusive) | + ### flow Fetches an flow given by its ID diff --git a/docs/graphql/object/usagebucket.md b/docs/graphql/object/usagebucket.md index a10b68d3..db502250 100644 --- a/docs/graphql/object/usagebucket.md +++ b/docs/graphql/object/usagebucket.md @@ -2,7 +2,7 @@ title: UsageBucket --- -Aggregated execution usage for a single day, week or month bucket +Aggregated usage for a single day, week or month bucket ## Fields without arguments @@ -10,5 +10,5 @@ Aggregated execution usage for a single day, week or month bucket |------|------|-------------| | `periodEnd` | [`ISO8601Date!`](../scalar/iso8601date.md) | End date of this usage bucket (inclusive) | | `periodStart` | [`ISO8601Date!`](../scalar/iso8601date.md) | Start date of this usage bucket (inclusive) | -| `usage` | [`BigInt!`](../scalar/bigint.md) | Number of executions in this bucket | -| `value` | [`Float!`](../scalar/float.md) | Total execution time in this bucket, in seconds | +| `usage` | [`BigInt!`](../scalar/bigint.md) | Number of events (e.g. executions, ai generations) in this bucket | +| `value` | [`Float!`](../scalar/float.md) | Aggregated value for this bucket (e.g. execution time in microseconds, ai usage in tokens) | diff --git a/extensions/cloud/spec/graphql/types/cloud/types/namespace_type_spec.rb b/extensions/cloud/spec/graphql/types/cloud/types/namespace_type_spec.rb index 42c10d04..c1cdc303 100644 --- a/extensions/cloud/spec/graphql/types/cloud/types/namespace_type_spec.rb +++ b/extensions/cloud/spec/graphql/types/cloud/types/namespace_type_spec.rb @@ -17,6 +17,7 @@ licenses currentLicense runtimeUsage + aiUsage userAbilities ] end diff --git a/extensions/ee/spec/graphql/ee/types/application_type_spec.rb b/extensions/ee/spec/graphql/ee/types/application_type_spec.rb index efd36837..2133dec9 100644 --- a/extensions/ee/spec/graphql/ee/types/application_type_spec.rb +++ b/extensions/ee/spec/graphql/ee/types/application_type_spec.rb @@ -15,6 +15,7 @@ identityProviders identityProviderLoginUrl runtimeUsage + aiUsage user_abilities ] end diff --git a/spec/factories/ai_usage_daily_aggregates.rb b/spec/factories/ai_usage_daily_aggregates.rb new file mode 100644 index 00000000..d98f3393 --- /dev/null +++ b/spec/factories/ai_usage_daily_aggregates.rb @@ -0,0 +1,12 @@ +# frozen_string_literal: true + +FactoryBot.define do + factory :ai_usage_daily_aggregate do + date { Time.zone.today } + flow_id { AiUsageDailyAggregate::NO_FLOW } + generation_count { 0 } + total_usage { 0 } + project { association(:namespace_project) } + namespace { project.namespace } + end +end diff --git a/spec/finders/ai_usage_daily_aggregates_finder_spec.rb b/spec/finders/ai_usage_daily_aggregates_finder_spec.rb new file mode 100644 index 00000000..878045b9 --- /dev/null +++ b/spec/finders/ai_usage_daily_aggregates_finder_spec.rb @@ -0,0 +1,76 @@ +# frozen_string_literal: true + +require 'rails_helper' + +RSpec.describe AiUsageDailyAggregatesFinder do + let(:project) { create(:namespace_project) } + let(:relation) { AiUsageDailyAggregate.where(project_id: project.id) } + + def seed_day(date, generation_count:, total_usage:) + # rubocop:disable Rails/SkipsModelValidations -- seeding pre-aggregated rows directly, not exercising validations + AiUsageDailyAggregate.insert!( + { + project_id: project.id, + flow_id: AiUsageDailyAggregate::NO_FLOW, + namespace_id: project.namespace_id, + date: date, + generation_count: generation_count, + total_usage: total_usage, + created_at: Time.current, + updated_at: Time.current, + } + ) + # rubocop:enable Rails/SkipsModelValidations + end + + describe '#execute' do + context 'with day aggregation' do + it 'returns one bucket per day in range' do + today = Time.zone.today + seed_day(today, generation_count: 2, total_usage: 200) + seed_day(today + 1.day, generation_count: 3, total_usage: 300) + + finder = described_class.new( + relation: relation, aggregation: 'day', after_date: today, before_date: today + 6.days + ) + + buckets = finder.execute + + expect(buckets.size).to eq(2) + expect(buckets.first).to have_attributes( + period_start: today, + period_end: today, + usage: 2, + value: 200 + ) + end + end + + context 'with month aggregation' do + it 'sums daily rows into a single monthly bucket' do + # Partitions only exist within the strategy's retention/headroom window (see + # Code0::ZeroTrack::Database::Partitioning::Strategy::Time), so dates must stay + # relative to today rather than hardcoded, or this test would eventually fail once + # the fixed dates fall outside the currently materialized partitions. + month_start = 1.month.ago.to_date.beginning_of_month + seed_day(month_start, generation_count: 1, total_usage: 100) + seed_day(month_start + 14.days, generation_count: 4, total_usage: 400) + + finder = described_class.new( + relation: relation, aggregation: 'month', after_date: month_start - 1.month, + before_date: month_start.end_of_month + ) + + buckets = finder.execute + + expect(buckets.size).to eq(1) + expect(buckets.first).to have_attributes( + period_start: month_start, + period_end: month_start.end_of_month, + usage: 5, + value: 500 + ) + end + end + end +end diff --git a/spec/finders/runtime_usage_daily_aggregates_finder_spec.rb b/spec/finders/runtime_usage_daily_aggregates_finder_spec.rb index 4d49dfce..dbd4a4cc 100644 --- a/spec/finders/runtime_usage_daily_aggregates_finder_spec.rb +++ b/spec/finders/runtime_usage_daily_aggregates_finder_spec.rb @@ -26,19 +26,20 @@ def seed_day(date, execution_count:, total_execution_time_us:) describe '#execute' do context 'with day aggregation' do it 'returns one bucket per day in range' do - seed_day(Date.new(2026, 8, 1), execution_count: 2, total_execution_time_us: 2_000_000) - seed_day(Date.new(2026, 8, 2), execution_count: 3, total_execution_time_us: 3_000_000) + today = Time.zone.today + seed_day(today, execution_count: 2, total_execution_time_us: 2_000_000) + seed_day(today + 1.day, execution_count: 3, total_execution_time_us: 3_000_000) finder = described_class.new( - relation: relation, aggregation: 'day', after_date: Date.new(2026, 8, 1), before_date: Date.new(2026, 8, 7) + relation: relation, aggregation: 'day', after_date: today, before_date: today + 6.days ) buckets = finder.execute expect(buckets.size).to eq(2) expect(buckets.first).to have_attributes( - period_start: Date.new(2026, 8, 1), - period_end: Date.new(2026, 8, 1), + period_start: today, + period_end: today, usage: 2, value: 2.0 ) @@ -47,20 +48,25 @@ def seed_day(date, execution_count:, total_execution_time_us:) context 'with month aggregation' do it 'sums daily rows into a single monthly bucket' do - seed_day(Date.new(2026, 6, 1), execution_count: 1, total_execution_time_us: 1_000_000) - seed_day(Date.new(2026, 6, 15), execution_count: 4, total_execution_time_us: 4_000_000) + # Partitions only exist within the strategy's retention/headroom window (see + # Code0::ZeroTrack::Database::Partitioning::Strategy::Time), so dates must stay + # relative to today rather than hardcoded, or this test would eventually fail once + # the fixed dates fall outside the currently materialized partitions. + month_start = 1.month.ago.to_date.beginning_of_month + seed_day(month_start, execution_count: 1, total_execution_time_us: 1_000_000) + seed_day(month_start + 14.days, execution_count: 4, total_execution_time_us: 4_000_000) finder = described_class.new( - relation: relation, aggregation: 'month', after_date: Date.new(2026, 5, 1), - before_date: Date.new(2026, 6, 30) + relation: relation, aggregation: 'month', after_date: month_start - 1.month, + before_date: month_start.end_of_month ) buckets = finder.execute expect(buckets.size).to eq(1) expect(buckets.first).to have_attributes( - period_start: Date.new(2026, 6, 1), - period_end: Date.new(2026, 6, 30), + period_start: month_start, + period_end: month_start.end_of_month, usage: 5, value: 5.0 ) diff --git a/spec/graphql/types/application_type_spec.rb b/spec/graphql/types/application_type_spec.rb index e7a24f16..b92e720c 100644 --- a/spec/graphql/types/application_type_spec.rb +++ b/spec/graphql/types/application_type_spec.rb @@ -13,6 +13,7 @@ identityProviders identityProviderLoginUrl runtimeUsage + aiUsage user_abilities ] end diff --git a/spec/graphql/types/concerns/has_runtime_usage_field_spec.rb b/spec/graphql/types/concerns/validates_usage_date_range_spec.rb similarity index 94% rename from spec/graphql/types/concerns/has_runtime_usage_field_spec.rb rename to spec/graphql/types/concerns/validates_usage_date_range_spec.rb index d1b34f46..859b3989 100644 --- a/spec/graphql/types/concerns/has_runtime_usage_field_spec.rb +++ b/spec/graphql/types/concerns/validates_usage_date_range_spec.rb @@ -2,7 +2,7 @@ require 'rails_helper' -RSpec.describe Types::Concerns::HasRuntimeUsageField do +RSpec.describe Types::Concerns::ValidatesUsageDateRange do describe '.validate_range!' do it 'raises when the range is shorter than the allowed 7-31 day span' do expect do diff --git a/spec/graphql/types/flow_spec.rb b/spec/graphql/types/flow_spec.rb index d7e44682..31bf7d87 100644 --- a/spec/graphql/types/flow_spec.rb +++ b/spec/graphql/types/flow_spec.rb @@ -19,6 +19,7 @@ execution_results linked_data_types runtime_usage + ai_usage user_abilities id created_at diff --git a/spec/graphql/types/namespace_project_type_spec.rb b/spec/graphql/types/namespace_project_type_spec.rb index f5b1a80c..50691f72 100644 --- a/spec/graphql/types/namespace_project_type_spec.rb +++ b/spec/graphql/types/namespace_project_type_spec.rb @@ -17,6 +17,7 @@ flows flow runtime_usage + ai_usage user_abilities created_at updated_at diff --git a/spec/graphql/types/namespace_type_spec.rb b/spec/graphql/types/namespace_type_spec.rb index a4a3e4ef..5ddf4908 100644 --- a/spec/graphql/types/namespace_type_spec.rb +++ b/spec/graphql/types/namespace_type_spec.rb @@ -13,6 +13,7 @@ project projects runtimeUsage + aiUsage userAbilities createdAt updatedAt diff --git a/spec/models/ai_usage_daily_aggregate_spec.rb b/spec/models/ai_usage_daily_aggregate_spec.rb new file mode 100644 index 00000000..c840ac5f --- /dev/null +++ b/spec/models/ai_usage_daily_aggregate_spec.rb @@ -0,0 +1,68 @@ +# frozen_string_literal: true + +require 'rails_helper' + +RSpec.describe AiUsageDailyAggregate do + describe 'associations' do + it { is_expected.to belong_to(:flow).inverse_of(:ai_usage_daily_aggregates).optional } + + it { + is_expected.to belong_to(:project).class_name('NamespaceProject') + .inverse_of(:ai_usage_daily_aggregates).optional + } + + it { is_expected.to belong_to(:namespace).inverse_of(:ai_usage_daily_aggregates).optional } + end + + describe 'validations' do + subject { create(:ai_usage_daily_aggregate) } + + it { is_expected.to validate_uniqueness_of(:date).scoped_to(%i[project_id flow_id]) } + end + + describe '.record_generation!' do + let(:project) { create(:namespace_project) } + let(:date) { Time.zone.today } + + def record(usage:, flow_id: described_class::NO_FLOW, record_date: date) + described_class.record_generation!( + project_id: project.id, namespace_id: project.namespace_id, flow_id: flow_id, + date: record_date, usage: usage, unique_by: %i[project_id flow_id date] + ) + end + + it 'creates a row on first use, denormalizing namespace_id' do + record(usage: 150) + + aggregate = described_class.find_by(project_id: project.id, flow_id: described_class::NO_FLOW, date: date) + expect(aggregate.generation_count).to eq(1) + expect(aggregate.total_usage).to eq(150) + expect(aggregate.namespace_id).to eq(project.namespace_id) + end + + it 'atomically accumulates on subsequent calls for the same project/flow/day' do + record(usage: 100) + record(usage: 200) + + aggregate = described_class.find_by(project_id: project.id, flow_id: described_class::NO_FLOW, date: date) + expect(aggregate.generation_count).to eq(2) + expect(aggregate.total_usage).to eq(300) + end + + it 'keeps separate rows for different flows in the same project/day' do + flow = create(:flow, project: project) + + record(usage: 100) + record(usage: 200, flow_id: flow.id) + + expect(described_class.where(project_id: project.id).count).to eq(2) + end + + it 'keeps separate rows for different days' do + record(usage: 100) + record(usage: 50, record_date: date - 1.day) + + expect(described_class.where(project_id: project.id).count).to eq(2) + end + end +end diff --git a/spec/services/velorum/generate_flow_service_spec.rb b/spec/services/velorum/generate_flow_service_spec.rb index 594867d9..1b086571 100644 --- a/spec/services/velorum/generate_flow_service_spec.rb +++ b/spec/services/velorum/generate_flow_service_spec.rb @@ -17,7 +17,7 @@ end let(:current_authentication) { instance_double(UserSession) } - let(:project) { instance_double(NamespaceProject, id: 12, primary_runtime: runtime) } + let(:project) { instance_double(NamespaceProject, id: 12, namespace_id: 3, primary_runtime: runtime) } let(:runtime) do instance_double( Runtime, @@ -133,6 +133,13 @@ ) end + it 'records the generation usage against the project with the no-flow sentinel' do + service_response + + aggregate = AiUsageDailyAggregate.find_by(project_id: project.id, flow_id: AiUsageDailyAggregate::NO_FLOW) + expect(aggregate).to have_attributes(generation_count: 1, total_usage: 42, namespace_id: project.namespace_id) + end + it 'omits definitions while the Velorum cache marker is still valid' do service_response @@ -157,6 +164,7 @@ let(:flow) do instance_double( Flow, + id: 55, project: project, to_generation_grpc: Tucana::Shared::GenerationFlow.new(name: 'Existing flow') ) @@ -171,6 +179,13 @@ expect(request.flow.name).to eq('Existing flow') end end + + it 'records the generation usage against the flow' do + service_response + + aggregate = AiUsageDailyAggregate.find_by(project_id: project.id, flow_id: flow.id) + expect(aggregate).to have_attributes(generation_count: 1, total_usage: 42, namespace_id: project.namespace_id) + end end context 'when the project does not have a primary runtime' do diff --git a/spec/services/velorum/record_generation_usage_service_spec.rb b/spec/services/velorum/record_generation_usage_service_spec.rb new file mode 100644 index 00000000..68aee6cf --- /dev/null +++ b/spec/services/velorum/record_generation_usage_service_spec.rb @@ -0,0 +1,44 @@ +# frozen_string_literal: true + +require 'rails_helper' + +RSpec.describe Velorum::RecordGenerationUsageService do + describe '#execute' do + let(:project) { create(:namespace_project) } + + context 'when no flow is given' do + it 'records a denormalized usage row keyed by the no-flow sentinel' do + described_class.new(project: project, usage: 42).execute + + aggregate = AiUsageDailyAggregate.find_by(project_id: project.id, flow_id: AiUsageDailyAggregate::NO_FLOW) + expect(aggregate).to have_attributes( + generation_count: 1, + total_usage: 42, + namespace_id: project.namespace_id + ) + expect(AiUsageDailyAggregate.count).to eq(1) + end + end + + context 'when a flow is given' do + let(:flow) { create(:flow, project: project) } + + it 'records a denormalized usage row for the flow/project/namespace' do + described_class.new(project: project, usage: 17, flow: flow).execute + + aggregate = AiUsageDailyAggregate.find_by(project_id: project.id, flow_id: flow.id) + expect(aggregate).to have_attributes( + generation_count: 1, + total_usage: 17, + namespace_id: project.namespace_id + ) + end + end + + it 'returns a successful service response' do + response = described_class.new(project: project, usage: 5).execute + + expect(response).to be_success + end + end +end