diff --git a/api/alembic/versions/2159d4ac431a_added_quota_tables.py b/api/alembic/versions/2159d4ac431a_added_quota_tables.py index 51efc4cc..24326e4b 100644 --- a/api/alembic/versions/2159d4ac431a_added_quota_tables.py +++ b/api/alembic/versions/2159d4ac431a_added_quota_tables.py @@ -18,6 +18,9 @@ branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +DEPRECATED_QUOTA_COMMENT = "Deprecated. MPS owns quota and credit ledger state." + + def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### # 1) Create the `quota_type` enum *before* we add the column that references it. @@ -34,7 +37,12 @@ def upgrade() -> None: sa.Column("organization_id", sa.Integer(), nullable=False), sa.Column("period_start", sa.DateTime(), nullable=False), sa.Column("period_end", sa.DateTime(), nullable=False), - sa.Column("quota_dograh_tokens", sa.Integer(), nullable=False), + sa.Column( + "quota_dograh_tokens", + sa.Integer(), + nullable=False, + comment=DEPRECATED_QUOTA_COMMENT, + ), sa.Column("used_dograh_tokens", sa.Integer(), nullable=False), sa.Column("created_at", sa.DateTime(timezone=True), nullable=True), sa.Column("updated_at", sa.DateTime(timezone=True), nullable=True), @@ -63,7 +71,11 @@ def upgrade() -> None: op.add_column( "organizations", sa.Column( - "quota_type", quota_type_enum, nullable=False, server_default="monthly" + "quota_type", + quota_type_enum, + nullable=False, + server_default="monthly", + comment=DEPRECATED_QUOTA_COMMENT, ), ) op.add_column( @@ -73,6 +85,7 @@ def upgrade() -> None: sa.Integer(), nullable=False, server_default=sa.text("0"), + comment=DEPRECATED_QUOTA_COMMENT, ), ) op.add_column( @@ -82,10 +95,17 @@ def upgrade() -> None: sa.Integer(), nullable=False, server_default=sa.text("LEAST(EXTRACT(DAY FROM CURRENT_DATE)::int, 28)"), + comment=DEPRECATED_QUOTA_COMMENT, ), ) op.add_column( - "organizations", sa.Column("quota_start_date", sa.DateTime(), nullable=True) + "organizations", + sa.Column( + "quota_start_date", + sa.DateTime(), + nullable=True, + comment=DEPRECATED_QUOTA_COMMENT, + ), ) op.add_column( "organizations", @@ -94,6 +114,7 @@ def upgrade() -> None: sa.Boolean(), nullable=False, server_default=sa.text("false"), + comment=DEPRECATED_QUOTA_COMMENT, ), ) # ### end Alembic commands ### diff --git a/api/alembic/versions/c425d3445750_add_columns_in_usage_table.py b/api/alembic/versions/c425d3445750_add_columns_in_usage_table.py index 998e7123..cbd9c654 100644 --- a/api/alembic/versions/c425d3445750_add_columns_in_usage_table.py +++ b/api/alembic/versions/c425d3445750_add_columns_in_usage_table.py @@ -18,6 +18,9 @@ branch_labels: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None +DEPRECATED_QUOTA_COMMENT = "Deprecated. MPS owns quota and credit ledger state." + + def upgrade() -> None: # ### commands auto generated by Alembic - please adjust! ### op.add_column( @@ -26,7 +29,12 @@ def upgrade() -> None: ) op.add_column( "organization_usage_cycles", - sa.Column("quota_amount_usd", sa.Float(), nullable=True), + sa.Column( + "quota_amount_usd", + sa.Float(), + nullable=True, + comment=DEPRECATED_QUOTA_COMMENT, + ), ) # ### end Alembic commands ### diff --git a/api/db/db_client.py b/api/db/db_client.py index de98cf19..15d1c108 100644 --- a/api/db/db_client.py +++ b/api/db/db_client.py @@ -53,7 +53,7 @@ class DBClient( - UserClient: handles user and user configuration operations - OrganizationClient: handles organization operations - OrganizationConfigurationClient: handles organization configuration operations - - OrganizationUsageClient: handles organization usage and quota operations + - OrganizationUsageClient: handles organization usage reporting aggregates - IntegrationClient: handles integration operations - WorkflowTemplateClient: handles workflow template operations - CampaignClient: handles campaign operations diff --git a/api/db/models.py b/api/db/models.py index c61cb03d..696cb6e6 100644 --- a/api/db/models.py +++ b/api/db/models.py @@ -97,22 +97,44 @@ class OrganizationModel(Base): provider_id = Column(String, unique=True, index=True, nullable=False) created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) - # Quota fields + # Deprecated: MPS owns quota and credit ledger state. quota_type = Column( Enum("monthly", "annual", name="quota_type"), nullable=False, default="monthly", server_default=text("'monthly'::quota_type"), + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, ) quota_dograh_tokens = Column( - Integer, nullable=False, default=0, server_default=text("0") + Integer, + nullable=False, + default=0, + server_default=text("0"), + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, ) quota_reset_day = Column( - Integer, nullable=False, default=1, server_default=text("1") - ) # 1-28, only for monthly - quota_start_date = Column(DateTime(timezone=True), nullable=True) # Only for annual + Integer, + nullable=False, + default=1, + server_default=text("1"), + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, + ) + quota_start_date = Column( + DateTime(timezone=True), + nullable=True, + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, + ) quota_enabled = Column( - Boolean, nullable=False, default=False, server_default=text("false") + Boolean, + nullable=False, + default=False, + server_default=text("false"), + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, ) price_per_second_usd = Column(Float, nullable=True) @@ -593,8 +615,9 @@ class WorkflowRunTextSessionModel(Base): class OrganizationUsageCycleModel(Base): """ - This model is used to track the usage of Dograh tokens for an organization for a given usage - cycle. + This model is used to track reporting aggregates for an organization for a given + usage cycle. Quota fields on this model are deprecated; MPS owns quota and + credit ledger state. """ __tablename__ = "organization_usage_cycles" @@ -603,14 +626,24 @@ class OrganizationUsageCycleModel(Base): organization_id = Column(Integer, ForeignKey("organizations.id"), nullable=False) period_start = Column(DateTime(timezone=True), nullable=False) period_end = Column(DateTime(timezone=True), nullable=False) - quota_dograh_tokens = Column(Integer, nullable=False) + quota_dograh_tokens = Column( + Integer, + nullable=False, + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, + ) used_dograh_tokens = Column(Float, nullable=False, default=0) total_duration_seconds = Column( Integer, nullable=False, default=0, server_default=text("0") ) # New USD tracking fields used_amount_usd = Column(Float, nullable=True, default=0) - quota_amount_usd = Column(Float, nullable=True) + quota_amount_usd = Column( + Float, + nullable=True, + comment="Deprecated. MPS owns quota and credit ledger state.", + info={"deprecated": True}, + ) created_at = Column(DateTime(timezone=True), default=lambda: datetime.now(UTC)) updated_at = Column( DateTime(timezone=True), diff --git a/api/db/organization_usage_client.py b/api/db/organization_usage_client.py index 826275e6..b147a2d6 100644 --- a/api/db/organization_usage_client.py +++ b/api/db/organization_usage_client.py @@ -23,7 +23,7 @@ from api.schemas.ai_model_configuration import EffectiveAIModelConfiguration class OrganizationUsageClient(BaseDBClient): - """Client for managing organization usage and quota operations.""" + """Client for managing organization usage reporting aggregates.""" async def get_or_create_current_cycle( self, organization_id: int, session=None @@ -49,14 +49,7 @@ class OrganizationUsageClient(BaseDBClient): self, organization_id: int, session, commit: bool ) -> OrganizationUsageCycleModel: """Internal implementation for get_or_create_current_cycle.""" - # Get organization to determine quota type - org_result = await session.execute( - select(OrganizationModel).where(OrganizationModel.id == organization_id) - ) - org = org_result.scalar_one() - - # Calculate current period - period_start, period_end = self._calculate_current_period(org) + period_start, period_end = self._calculate_current_period() # Try to get existing cycle cycle_result = await session.execute( @@ -78,7 +71,8 @@ class OrganizationUsageClient(BaseDBClient): organization_id=organization_id, period_start=period_start, period_end=period_end, - quota_dograh_tokens=org.quota_dograh_tokens, + # Deprecated non-null column retained for historical schema compatibility. + quota_dograh_tokens=0, ) # Handle concurrent inserts gracefully stmt = stmt.on_conflict_do_nothing( @@ -102,54 +96,6 @@ class OrganizationUsageClient(BaseDBClient): ) return cycle_result.scalar_one() - async def check_and_reserve_quota( - self, organization_id: int, estimated_tokens: int = 0 - ) -> bool: - """ - Check if organization has sufficient quota and optionally reserve tokens. - Returns True if quota is available, False otherwise. - - This method is fully atomic and safe for concurrent access from multiple processes. - """ - async with self.async_session() as session: - # Get organization - org_result = await session.execute( - select(OrganizationModel).where(OrganizationModel.id == organization_id) - ) - org = org_result.scalar_one_or_none() - - if not org or not org.quota_enabled: - # No quota enforcement if not enabled - return True - - # Get or create current cycle within the same session/transaction - cycle = await self._get_or_create_current_cycle_impl( - organization_id, session, commit=False - ) - - # Atomic check and update with row-level lock - result = await session.execute( - select(OrganizationUsageCycleModel) - .where( - and_( - OrganizationUsageCycleModel.id == cycle.id, - OrganizationUsageCycleModel.used_dograh_tokens - + estimated_tokens - <= OrganizationUsageCycleModel.quota_dograh_tokens, - ) - ) - .with_for_update(skip_locked=False) - ) - - cycle_locked = result.scalar_one_or_none() - if cycle_locked: - # Update the usage atomically - cycle_locked.used_dograh_tokens += estimated_tokens - await session.commit() - return True - - return False - async def update_usage_after_run( self, organization_id: int, @@ -188,9 +134,8 @@ class OrganizationUsageClient(BaseDBClient): await session.commit() async def get_current_usage(self, organization_id: int) -> dict: - """Get current period usage information.""" + """Get current reporting-period usage information.""" async with self.async_session() as session: - # Get organization org_result = await session.execute( select(OrganizationModel).where(OrganizationModel.id == organization_id) ) @@ -201,42 +146,19 @@ class OrganizationUsageClient(BaseDBClient): organization_id, session, commit=False ) - # Calculate next refresh date - if org.quota_type == "monthly": - next_refresh = cycle.period_end + relativedelta(days=1) - else: # annual - next_refresh = cycle.period_end + relativedelta(days=1) - result = { "period_start": cycle.period_start.isoformat(), "period_end": cycle.period_end.isoformat(), "used_dograh_tokens": cycle.used_dograh_tokens, - "quota_dograh_tokens": cycle.quota_dograh_tokens, - "percentage_used": ( - round( - (cycle.used_dograh_tokens / cycle.quota_dograh_tokens) * 100, 2 - ) - if cycle.quota_dograh_tokens > 0 - else 0 - ), - "next_refresh_date": next_refresh.date().isoformat(), - "quota_enabled": org.quota_enabled, "total_duration_seconds": cycle.total_duration_seconds, } # Add USD fields if organization has pricing if org.price_per_second_usd is not None: result["used_amount_usd"] = cycle.used_amount_usd or 0 - result["quota_amount_usd"] = cycle.quota_amount_usd result["currency"] = "USD" result["price_per_second_usd"] = org.price_per_second_usd - # Calculate percentage based on USD if available - if cycle.quota_amount_usd and cycle.quota_amount_usd > 0: - result["percentage_used"] = round( - ((cycle.used_amount_usd or 0) / cycle.quota_amount_usd) * 100, 2 - ) - return result async def get_usage_history( @@ -545,83 +467,11 @@ class OrganizationUsageClient(BaseDBClient): "currency": "USD", } - async def update_organization_quota( - self, - organization_id: int, - quota_type: str, - quota_dograh_tokens: int, - quota_reset_day: Optional[int] = None, - quota_start_date: Optional[datetime] = None, - ) -> OrganizationModel: - """Update organization quota settings.""" - async with self.async_session() as session: - result = await session.execute( - select(OrganizationModel).where(OrganizationModel.id == organization_id) - ) - org = result.scalar_one() - - org.quota_type = quota_type - org.quota_dograh_tokens = quota_dograh_tokens - org.quota_enabled = True - - if quota_type == "monthly" and quota_reset_day: - org.quota_reset_day = quota_reset_day - elif quota_type == "annual" and quota_start_date: - org.quota_start_date = quota_start_date - - await session.commit() - await session.refresh(org) - return org - - def _calculate_current_period( - self, org: OrganizationModel - ) -> tuple[datetime, datetime]: - """Calculate the current billing period based on organization settings.""" + def _calculate_current_period(self) -> tuple[datetime, datetime]: + """Calculate the current calendar-month reporting period.""" now = datetime.now(timezone.utc) - if org.quota_type == "monthly": - # Find the start of the current billing month - reset_day = org.quota_reset_day - - # Handle month boundaries - if now.day >= reset_day: - period_start = now.replace( - day=reset_day, hour=0, minute=0, second=0, microsecond=0 - ) - else: - # Previous month - period_start = (now - relativedelta(months=1)).replace( - day=reset_day, hour=0, minute=0, second=0, microsecond=0 - ) - - # End is one month later minus 1 second - period_end = ( - period_start + relativedelta(months=1) - relativedelta(seconds=1) - ) - - else: # annual - if not org.quota_start_date: - # Default to calendar year - period_start = now.replace( - month=1, day=1, hour=0, minute=0, second=0, microsecond=0 - ) - period_end = ( - period_start + relativedelta(years=1) - relativedelta(seconds=1) - ) - else: - # Find current annual period - start_date = org.quota_start_date.replace(tzinfo=timezone.utc) - years_diff = now.year - start_date.year - - # Adjust for whether we've passed the anniversary - if now.month < start_date.month or ( - now.month == start_date.month and now.day < start_date.day - ): - years_diff -= 1 - - period_start = start_date + relativedelta(years=years_diff) - period_end = ( - period_start + relativedelta(years=1) - relativedelta(seconds=1) - ) + period_start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + period_end = period_start + relativedelta(months=1) - relativedelta(seconds=1) return period_start, period_end diff --git a/api/routes/organization_usage.py b/api/routes/organization_usage.py index 64a84956..1ce76acf 100644 --- a/api/routes/organization_usage.py +++ b/api/routes/organization_usage.py @@ -25,14 +25,8 @@ class CurrentUsageResponse(BaseModel): period_start: str period_end: str used_dograh_tokens: float - quota_dograh_tokens: int - percentage_used: float - next_refresh_date: str - quota_enabled: bool total_duration_seconds: int - # New USD fields used_amount_usd: Optional[float] = None - quota_amount_usd: Optional[float] = None currency: Optional[str] = None price_per_second_usd: Optional[float] = None @@ -134,7 +128,7 @@ class DailyUsageBreakdownResponse(BaseModel): @router.get("/usage/current-period", response_model=CurrentUsageResponse) async def get_current_period_usage(user: UserModel = Depends(get_user)): - """Get current billing period usage for the user's organization.""" + """Get current reporting-period usage for the user's organization.""" if not user.selected_organization_id: raise HTTPException(status_code=400, detail="No organization selected") diff --git a/ui/src/app/reports/page.tsx b/ui/src/app/reports/page.tsx index debee16b..3f703bdc 100644 --- a/ui/src/app/reports/page.tsx +++ b/ui/src/app/reports/page.tsx @@ -1,14 +1,12 @@ 'use client'; import { addDays, format, subDays } from 'date-fns'; -import { Calendar, ChevronLeft, ChevronRight, CreditCard, Download, Loader2 } from 'lucide-react'; +import { Calendar, ChevronLeft, ChevronRight, Download } from 'lucide-react'; import { useEffect, useState } from 'react'; -import { client } from '@/client/client.gen'; import { getDailyReportApiV1OrganizationsReportsDailyGet, getDailyRunsDetailApiV1OrganizationsReportsDailyRunsGet, - getModelConfigurationV2ApiV1OrganizationsModelConfigurationsV2Get, getPreferencesApiV1OrganizationsPreferencesGet, getWorkflowOptionsApiV1OrganizationsReportsWorkflowsGet } from '@/client/sdk.gen'; @@ -19,7 +17,6 @@ import { Card } from '@/components/ui/card'; import { Popover, PopoverContent, PopoverTrigger } from '@/components/ui/popover'; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from '@/components/ui/select'; import { Skeleton } from '@/components/ui/skeleton'; -import { detailFromError } from '@/lib/apiError'; import { useAuth } from '@/lib/auth'; import { DispositionChart } from './components/DispositionChart'; @@ -53,10 +50,6 @@ interface DailyReport { }>; } -type CreditPurchaseUrlResponse = { - checkout_url: string; -}; - export default function ReportsPage() { const [selectedDate, setSelectedDate] = useState(new Date()); const [selectedWorkflow, setSelectedWorkflow] = useState('all'); @@ -64,9 +57,6 @@ export default function ReportsPage() { const [report, setReport] = useState(null); const [loading, setLoading] = useState(true); const [error, setError] = useState(null); - const [canBuyCredits, setCanBuyCredits] = useState(false); - const [purchaseLoading, setPurchaseLoading] = useState(false); - const [purchaseError, setPurchaseError] = useState(null); const [timezone, setTimezone] = useState('America/New_York'); const auth = useAuth(); @@ -104,50 +94,6 @@ export default function ReportsPage() { fetchPreferences(); }, [auth.isAuthenticated]); - useEffect(() => { - if (auth.loading) return; - - if (!auth.isAuthenticated) { - setCanBuyCredits(false); - return; - } - - let cancelled = false; - - const fetchCreditPurchaseAvailability = async () => { - try { - const response = await getModelConfigurationV2ApiV1OrganizationsModelConfigurationsV2Get(); - if (cancelled) return; - - if (response.error || !response.data) { - setCanBuyCredits(false); - return; - } - - const configuration = response.data.configuration as { - mode?: unknown; - dograh?: unknown; - } | null; - setCanBuyCredits( - response.data.source === 'organization_v2' && - configuration?.mode === 'dograh' && - Boolean(configuration.dograh) - ); - } catch (err) { - console.error('Failed to check credit purchase availability:', err); - if (!cancelled) { - setCanBuyCredits(false); - } - } - }; - - fetchCreditPurchaseAvailability(); - - return () => { - cancelled = true; - }; - }, [auth.loading, auth.isAuthenticated]); - // Fetch report data when date or workflow changes useEffect(() => { const fetchReport = async () => { @@ -249,41 +195,6 @@ export default function ReportsPage() { } }; - const handleBuyCredits = async () => { - if (!auth.isAuthenticated || purchaseLoading) return; - - setPurchaseLoading(true); - setPurchaseError(null); - - try { - const response = await client.post< - { 200: CreditPurchaseUrlResponse }, - { detail?: unknown } - >({ - url: '/api/v1/organizations/usage/mps-credits/purchase-url', - }); - - if (response.error) { - throw new Error( - detailFromError(response.error, 'Failed to create credit purchase URL') - ); - } - - const checkoutUrl = response.data?.checkout_url; - if (!checkoutUrl) { - throw new Error('Failed to create credit purchase URL'); - } - - window.location.href = checkoutUrl; - } catch (err) { - console.error('Failed to create credit purchase URL:', err); - setPurchaseError( - err instanceof Error ? err.message : 'Failed to create credit purchase URL' - ); - setPurchaseLoading(false); - } - }; - const isToday = format(selectedDate, 'yyyy-MM-dd') === format(new Date(), 'yyyy-MM-dd'); return ( @@ -291,27 +202,7 @@ export default function ReportsPage() { {/* Header */}
-
-

Daily Reports

- {canBuyCredits && ( - - )} -
- {purchaseError && ( -

{purchaseError}

- )} +

Daily Reports

{/* Date Navigation & Workflow Selector */} diff --git a/ui/src/client/sdk.gen.ts b/ui/src/client/sdk.gen.ts index 4ae08c98..c32811da 100644 --- a/ui/src/client/sdk.gen.ts +++ b/ui/src/client/sdk.gen.ts @@ -1239,7 +1239,7 @@ export const reactivateServiceKeyApiV1UserServiceKeysServiceKeyIdReactivatePut = /** * Get Current Period Usage * - * Get current billing period usage for the user's organization. + * Get current reporting-period usage for the user's organization. */ export const getCurrentPeriodUsageApiV1OrganizationsUsageCurrentPeriodGet = (options?: Options) => (options?.client ?? client).get({ url: '/api/v1/organizations/usage/current-period', ...options }); diff --git a/ui/src/client/types.gen.ts b/ui/src/client/types.gen.ts index 212e0d3f..a0dc5b72 100644 --- a/ui/src/client/types.gen.ts +++ b/ui/src/client/types.gen.ts @@ -1642,22 +1642,6 @@ export type CurrentUsageResponse = { * Used Dograh Tokens */ used_dograh_tokens: number; - /** - * Quota Dograh Tokens - */ - quota_dograh_tokens: number; - /** - * Percentage Used - */ - percentage_used: number; - /** - * Next Refresh Date - */ - next_refresh_date: string; - /** - * Quota Enabled - */ - quota_enabled: boolean; /** * Total Duration Seconds */ @@ -1666,10 +1650,6 @@ export type CurrentUsageResponse = { * Used Amount Usd */ used_amount_usd?: number | null; - /** - * Quota Amount Usd - */ - quota_amount_usd?: number | null; /** * Currency */