feat(subscriptions): gate AI subscriptions on AI credit budget#59634
feat(subscriptions): gate AI subscriptions on AI credit budget#59634vdekrijger wants to merge 1 commit into
Conversation
|
Warning This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
This stack of pull requests is managed by Graphite. Learn more about stacking. |
|
🎭 Playwright didn't run on this PR — your changes touch code that could affect E2E behavior, but Playwright is opt-in via label now to keep CI cost down. Add the Most PRs don't need this. Real regressions still get caught on master and fix-forward. |
🔍 Migration Risk AnalysisWe've analyzed your migrations for potential risks. Summary: 0 Safe | 0 Needs Review | 1 Blocked ❌ BlockedCauses locks or breaks compatibility 📚 How to Deploy These Changes SafelyAddField: This operation acquires a brief lock but doesn't rewrite the table. Deployment uses lock timeouts with automatic retries, so lock contention will cause retries rather than connection pile-up. RunPython: Use batching for large data migrations:
See the migration safety guide Last updated: 2026-05-27 08:06 UTC (cb7ecf6) |
| # adapted: ValidationError there → 403 Response here. Helper-driven so a | ||
| # new gate (e.g. quota) only needs adding in one place. | ||
| gate_reason = _ai_create_gate_reason(self.organization, kind="reports", verb="generating") | ||
| if gate_reason is not None: |
There was a problem hiding this comment.
[Re: line +1015]
iirc this function will be removed, so no need to make the changes here.
See this comment inline on Graphite.
There was a problem hiding this comment.
Similar to the test
There was a problem hiding this comment.
Reverted — removed the AI-credit gate from ai_report (and its test test_ai_report_endpoint_rejects_when_over_ai_credit_limit) plus the now-unused import, since this endpoint is going away.
| """Whether the team has exhausted its AI credit budget, per the billing quota-limiting cache — | ||
| the same signal the chat assistant enforces (ee/api/conversation.py). Shared by every | ||
| LLM-spending subscription path (AI reports, AI summaries) so the resource + cache-key pair lives |
There was a problem hiding this comment.
we can simplify this pydoc imo, lets just focus on the why! and we should probably also update the checks around the AI summaries in the "normal" (insight / dashboard) subscriptions
There was a problem hiding this comment.
Simplified the docstring to focus on the why. Also routed the AI-summary path (snapshot_activities.py) through the shared is_team_over_ai_credit_budget helper, and migrated the remaining AI-credit callers (conversation.py Max chat, subscription.py summary gate) so every LLM-spending path checks the same resource + cache-key pair.
| """Notify the subscription owner that a scheduled AI report was skipped because the org | ||
| is out of AI credits. `billing_period_key` keys the campaign so MessagingRecord dedups to | ||
| one notice per billing period even if the skip path runs more than once — this, not the | ||
| reschedule alone, is what guarantees we email at most once per credit-reset cycle.""" |
There was a problem hiding this comment.
We can simplify this pydoc as well imo.
There was a problem hiding this comment.
Simplified the docstring.
| super().setUp() | ||
| # An AI subscription can only exist for an org that approved AI data processing — creation | ||
| # is gated on it. Delivery re-checks consent at send time, so the suite must reflect the | ||
| # real precondition; otherwise every delivery short-circuits as `ai_consent_revoked`. |
There was a problem hiding this comment.
imo we can remove this pydoc!
There was a problem hiding this comment.
Removed the comment.
| @patch("posthog.temporal.subscriptions.activities.send_email_ai_subscription_credit_limited") | ||
| @patch("posthog.temporal.subscriptions.activities.generate_ai_subscription_markdown") | ||
| @patch("posthog.temporal.subscriptions.activities.is_team_over_ai_credit_budget") | ||
| def test_cached_markdown_delivers_even_when_over_credit_limit( |
There was a problem hiding this comment.
I think we want to reevaluate how we do this caching behaviour thing, so we might have to change this test.
There was a problem hiding this comment.
Left as-is for now — per your note the caching flow is being reworked downstack (#59631), so I didn't want to collide with that in-flight change. This test will get updated when that lands.
| """When the org's AI credits next reset — the billing-period end synced from billing.""" | ||
| usage = subscription.team.organization.usage | ||
| period = usage.get("period") if usage else None | ||
| if period and len(period) == 2 and period[1]: |
There was a problem hiding this comment.
How does the period model look like? What is on index 0?
There was a problem hiding this comment.
organization.usage["period"] is [current_period_start, current_period_end] as ISO strings (set in billing_manager.py). Index 0 is the period start, index 1 the end — which is when AI credits reset, so the code reads period[1]. Added a clarifying comment.
| Persisting `next_delivery_date = reset_date` is deliberate: the always-runs | ||
| `advance_next_delivery_date` activity recomputes from this value (`rrule.after(reset_date)`), | ||
| so the next delivery lands on the first on-schedule slot AFTER credits reset rather than the | ||
| next normal occurrence (which could fall before the reset and re-fire while still over-limit). |
There was a problem hiding this comment.
Same with this pydoc, please focus on the why
There was a problem hiding this comment.
Simplified — kept just the non-obvious why (persisting next_delivery_date = reset_date so the always-runs advance_next_delivery_date recomputes the next slot past the credit reset).
The ad-hoc ai_report endpoint and scheduled AI deliveries enforce the org's AI credit budget via a shared is_team_over_ai_credit_budget helper — the same billing signal the chat assistant (ee/api/conversation.py) and the AI-summary path use. - ai_report endpoint returns 402 when over budget, before spending tokens - scheduled deliveries fail open, skip + reschedule past the credit reset, notify the owner once per period, and increment a Prometheus counter (parity with the AI-summary skip metric). Cache hits still deliver — their tokens were already spent. - shared helper in ee/billing/quota_limiting keeps the resource + cache-key pair in one place so the AI-report and AI-summary paths can't drift
ce31c52 to
cb7ecf6
Compare

Problem
Changes
How did you test this code?
👉 Stay up-to-date with PostHog coding conventions for a smoother review.
Publish to changelog?
Docs update
🤖 Agent context