Skip to content

Commit c41f182

Browse files
adboioadboio
authored
refactor(billing): consolidate cutover logic per review pass
Move usage-limit modal copy derivation into core (pure, unit-tested) and thin the component; persist the billing announcement via the shared electronStorage-backed persist middleware with a hydration guard; gate seat fetch/provision once in the seat store instead of at three call sites; share a GLM-visibility hook and a models-cache validity helper; name the code_usage_billed tri-state with isCodeUsageUnbilled; merge the session service's duplicated limit branches. Generated-By: PostHog Code Task-Id: bfbcd747-a365-429f-a5dc-622e2c2f7edf
1 parent e11ad1a commit c41f182

17 files changed

Lines changed: 370 additions & 245 deletions

‎packages/agent/src/gateway-models.ts‎

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -78,12 +78,26 @@ const CACHE_TTL = 10 * 60 * 1000; // 10 minutes
7878
// the callers fall through to `return []`.
7979
const GATEWAY_FETCH_TIMEOUT_MS = 10_000;
8080

81-
let gatewayModelsCache: {
82-
models: GatewayModel[];
81+
// Authed and anonymous responses differ (free-tier marks are only present on
82+
// authed fetches), so cached entries are keyed on auth presence: a cached
83+
// anonymous list must not serve an authed caller or vice versa.
84+
interface ModelsCache<T> {
85+
models: T[];
8386
expiry: number;
8487
url: string;
8588
authed: boolean;
86-
} | null = null;
89+
}
90+
91+
function readModelsCache<T>(
92+
cache: ModelsCache<T> | null,
93+
url: string,
94+
authed: boolean,
95+
): T[] | null {
96+
if (!cache || cache.url !== url || cache.authed !== authed) return null;
97+
return Date.now() < cache.expiry ? cache.models : null;
98+
}
99+
100+
let gatewayModelsCache: ModelsCache<GatewayModel> | null = null;
87101

88102
function authHeaders(authToken?: string): Record<string, string> | undefined {
89103
return authToken ? { Authorization: `Bearer ${authToken}` } : undefined;
@@ -97,18 +111,9 @@ export async function fetchGatewayModels(
97111
return [];
98112
}
99113

100-
// Authed and anonymous responses differ (free-tier marks are only present
101-
// on authed fetches), so a cached anonymous list must not serve an authed
102-
// caller or vice versa.
103114
const authed = Boolean(options?.authToken);
104-
if (
105-
gatewayModelsCache &&
106-
gatewayModelsCache.url === gatewayUrl &&
107-
gatewayModelsCache.authed === authed &&
108-
Date.now() < gatewayModelsCache.expiry
109-
) {
110-
return gatewayModelsCache.models;
111-
}
115+
const cached = readModelsCache(gatewayModelsCache, gatewayUrl, authed);
116+
if (cached) return cached;
112117

113118
const modelsUrl = `${gatewayUrl}/v1/models`;
114119

@@ -175,12 +180,7 @@ export interface ModelInfo {
175180
restriction_reason?: string | null;
176181
}
177182

178-
let modelsListCache: {
179-
models: ModelInfo[];
180-
expiry: number;
181-
url: string;
182-
authed: boolean;
183-
} | null = null;
183+
let modelsListCache: ModelsCache<ModelInfo> | null = null;
184184

185185
export async function fetchModelsList(
186186
options?: FetchGatewayModelsOptions,
@@ -191,14 +191,8 @@ export async function fetchModelsList(
191191
}
192192

193193
const authed = Boolean(options?.authToken);
194-
if (
195-
modelsListCache &&
196-
modelsListCache.url === gatewayUrl &&
197-
modelsListCache.authed === authed &&
198-
Date.now() < modelsListCache.expiry
199-
) {
200-
return modelsListCache.models;
201-
}
194+
const cached = readModelsCache(modelsListCache, gatewayUrl, authed);
195+
if (cached) return cached;
202196

203197
try {
204198
const modelsUrl = `${gatewayUrl}/v1/models`;

‎packages/core/src/billing/usageDisplay.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,17 @@ export function isUsageExceeded(usage: UsageOutput): boolean {
99
);
1010
}
1111

12+
/**
13+
* The org is confirmed on the free tier (not billed for Code usage). False
14+
* when billed OR when the state is unknown — `code_usage_billed` is absent on
15+
* gateways predating the field, and absence must never read as free.
16+
*/
17+
export function isCodeUsageUnbilled(
18+
usage: Pick<UsageOutput, "code_usage_billed"> | null | undefined,
19+
): boolean {
20+
return usage?.code_usage_billed === false;
21+
}
22+
1223
export function formatResetTime(
1324
resetAtIso: string,
1425
now: number = Date.now(),

‎packages/core/src/billing/usageLimitContent.test.ts‎

Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
import { describe, expect, it } from "vitest";
2+
import {
3+
deriveUsageLimitCause,
4+
seatEraLimitContent,
5+
usageBasedLimitContent,
6+
} from "./usageLimitContent";
7+
8+
describe("deriveUsageLimitCause", () => {
9+
it.each([
10+
["model_gate", "burst", "model_gate"],
11+
[null, "burst", "user_daily_limit"],
12+
[null, "sustained", "user_monthly_limit"],
13+
[null, null, "org_limit"],
14+
] as const)("cause %s + bucket %s -> %s", (cause, bucket, expected) => {
15+
expect(deriveUsageLimitCause(cause, bucket)).toBe(expected);
16+
});
17+
});
18+
19+
describe("usageBasedLimitContent", () => {
20+
it("names the gated model and offers a payment method", () => {
21+
const content = usageBasedLimitContent({
22+
cause: "model_gate",
23+
model: "Claude Opus 4.8",
24+
resetLabel: null,
25+
billed: false,
26+
});
27+
expect(content.title).toBe("Unlock premium models");
28+
expect(content.description).toContain("Claude Opus 4.8 isn't");
29+
expect(content.actionLabel).toBe("Add payment method");
30+
});
31+
32+
it("falls back to generic wording when the gated model is unknown", () => {
33+
const content = usageBasedLimitContent({
34+
cause: "model_gate",
35+
model: null,
36+
resetLabel: null,
37+
billed: undefined,
38+
});
39+
expect(content.description).toContain("This model isn't");
40+
});
41+
42+
it.each([
43+
// Confirmed-free org: allocation used up, the fix is adding a card.
44+
[false, "Free usage used up", "Add payment method"],
45+
// Billed org: the fix is raising the spend limit.
46+
[true, "Organization usage limit reached", "Manage billing"],
47+
// Unknown billed state must not read as free.
48+
[undefined, "Organization usage limit reached", "Manage billing"],
49+
] as const)(
50+
"org_limit with billed=%s -> %s / %s",
51+
(billed, title, actionLabel) => {
52+
const content = usageBasedLimitContent({
53+
cause: "org_limit",
54+
model: null,
55+
resetLabel: null,
56+
billed,
57+
});
58+
expect(content.title).toBe(title);
59+
expect(content.actionLabel).toBe(actionLabel);
60+
},
61+
);
62+
63+
it.each([
64+
["user_daily_limit", "Free daily limit reached"],
65+
["user_monthly_limit", "Free monthly limit reached"],
66+
] as const)("%s carries the reset hint", (cause, title) => {
67+
const content = usageBasedLimitContent({
68+
cause,
69+
model: null,
70+
resetLabel: "Resets in 2h",
71+
billed: false,
72+
});
73+
expect(content.title).toBe(title);
74+
expect(content.description).toContain("Resets in 2h");
75+
expect(content.actionLabel).toBe("Add payment method");
76+
});
77+
});
78+
79+
describe("seatEraLimitContent", () => {
80+
it("keeps the Pro cap copy with no upgrade action", () => {
81+
const content = seatEraLimitContent({
82+
bucket: "sustained",
83+
isPro: true,
84+
resetLabel: "Resets in 3d",
85+
});
86+
expect(content.title).toBe("Monthly limit reached");
87+
expect(content.description).toContain("monthly usage cap");
88+
expect(content.actionLabel).toBeNull();
89+
expect(content.dismissLabel).toBe("Got it");
90+
});
91+
92+
it("keeps the free-plan upgrade pitch", () => {
93+
const content = seatEraLimitContent({
94+
bucket: "burst",
95+
isPro: false,
96+
resetLabel: null,
97+
});
98+
expect(content.title).toBe("Daily limit reached");
99+
expect(content.description).toContain("Upgrade to Pro for 40×");
100+
expect(content.actionLabel).toBe("See Pro");
101+
});
102+
});

‎packages/core/src/billing/usageLimitContent.ts‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
import type { GatewayLimitCause } from "@posthog/shared";
2+
import { PRO_USAGE_MULTIPLIER } from "./usageDisplay";
3+
4+
export interface UsageLimitContent {
5+
title: string;
6+
description: string;
7+
/** Primary action label; null renders only the dismiss button. */
8+
actionLabel: string | null;
9+
dismissLabel: string;
10+
}
11+
12+
/**
13+
* Maps a legacy bucket-only caller onto a usage-based cause: no exceeded
14+
* bucket means the org's credit bucket tripped the limit.
15+
*/
16+
export function deriveUsageLimitCause(
17+
cause: GatewayLimitCause | null,
18+
bucket: "burst" | "sustained" | null,
19+
): GatewayLimitCause {
20+
if (cause) return cause;
21+
if (bucket === "burst") return "user_daily_limit";
22+
if (bucket === "sustained") return "user_monthly_limit";
23+
return "org_limit";
24+
}
25+
26+
export function usageBasedLimitContent(args: {
27+
cause: GatewayLimitCause;
28+
model: string | null;
29+
resetLabel: string | null;
30+
/** usage.code_usage_billed — absent means unknown, not free. */
31+
billed: boolean | undefined;
32+
}): UsageLimitContent {
33+
const { cause, model, resetLabel, billed } = args;
34+
35+
if (cause === "model_gate") {
36+
return {
37+
title: "Unlock premium models",
38+
description: `${model ? `${model} isn't` : "This model isn't"} included in the free tier. Add a payment method to your organization to unlock all models — you only pay for what you use. You can keep working now by switching to an included model.`,
39+
actionLabel: "Add payment method",
40+
dismissLabel: "Not now",
41+
};
42+
}
43+
44+
if (cause === "org_limit") {
45+
if (billed === false) {
46+
return {
47+
title: "Free usage used up",
48+
description:
49+
"Your organization has used its included PostHog Code usage for this billing period. Add a payment method to keep going — you only pay for what you use.",
50+
actionLabel: "Add payment method",
51+
dismissLabel: "Not now",
52+
};
53+
}
54+
return {
55+
title: "Organization usage limit reached",
56+
description:
57+
"Your organization has reached its PostHog Code spend limit for this billing period. Raise or remove the limit in your PostHog billing settings to keep going.",
58+
actionLabel: "Manage billing",
59+
dismissLabel: "Got it",
60+
};
61+
}
62+
63+
const period = cause === "user_daily_limit" ? "daily" : "monthly";
64+
return {
65+
title: `Free ${period} limit reached`,
66+
description: `You've hit the free tier's ${period} usage limit.${
67+
resetLabel ? ` ${resetLabel}.` : ""
68+
} Add a payment method to your organization for uncapped usage-based access.`,
69+
actionLabel: "Add payment method",
70+
dismissLabel: "Not now",
71+
};
72+
}
73+
74+
export function seatEraLimitContent(args: {
75+
bucket: "burst" | "sustained" | null;
76+
isPro: boolean;
77+
resetLabel: string | null;
78+
}): UsageLimitContent {
79+
const { bucket, isPro, resetLabel } = args;
80+
const isDaily = bucket === "burst";
81+
const isMonthly = bucket === "sustained";
82+
83+
const title = isDaily
84+
? "Daily limit reached"
85+
: isMonthly && !isPro
86+
? "You're out of usage for this month"
87+
: isMonthly
88+
? "Monthly limit reached"
89+
: "Usage limit reached";
90+
91+
const proCapLabel = isDaily
92+
? "a daily usage cap"
93+
: isMonthly
94+
? "a monthly usage cap"
95+
: "usage caps";
96+
const description = isPro
97+
? `Your Pro plan has ${proCapLabel}.${resetLabel ? ` ${resetLabel}.` : ""}`
98+
: `You've hit your Free ${
99+
isDaily ? "daily" : isMonthly ? "monthly" : "usage"
100+
} limit. Upgrade to Pro for ${PRO_USAGE_MULTIPLIER}× more usage.`;
101+
102+
return {
103+
title,
104+
description,
105+
actionLabel: isPro ? null : "See Pro",
106+
dismissLabel: isPro ? "Got it" : "Not now",
107+
};
108+
}

‎packages/core/src/sessions/sessionService.ts‎

Lines changed: 20 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -2716,33 +2716,33 @@ export class SessionService {
27162716

27172717
const limitCause = classifyGatewayLimitError(errorMessage, errorDetails);
27182718

2719-
// Free-tier model gate (403): the session is healthy — the user needs
2720-
// to switch model or add a payment method. Surface the upgrade gate,
2721-
// not an error state.
2722-
if (limitCause === "model_gate") {
2723-
this.d.log.warn("Model gated for free tier, showing upgrade gate", {
2719+
// Gateway billing denials — free-tier model gate (403) or a usage
2720+
// limit (429): the session is healthy, the fix is switching model,
2721+
// adding a payment method, or raising a limit. Surface the upgrade
2722+
// gate, not an error state.
2723+
if (
2724+
limitCause === "model_gate" ||
2725+
isRateLimitError(errorMessage, errorDetails)
2726+
) {
2727+
this.d.log.warn("Gateway limit reached, showing usage limit modal", {
27242728
taskRunId: session.taskRunId,
2725-
});
2726-
this.d.store.updateSession(session.taskRunId, {
2727-
isPromptPending: false,
2728-
promptStartedAt: null,
2729-
});
2730-
this.d.usageLimit.show({
27312729
cause: limitCause,
2732-
model: extractGatedModel(errorMessage, errorDetails) ?? undefined,
2733-
});
2734-
return { stopReason: "rate_limited" };
2735-
}
2736-
2737-
if (isRateLimitError(errorMessage, errorDetails)) {
2738-
this.d.log.warn("Rate limit exceeded, showing usage limit modal", {
2739-
taskRunId: session.taskRunId,
27402730
});
27412731
this.d.store.updateSession(session.taskRunId, {
27422732
isPromptPending: false,
27432733
promptStartedAt: null,
27442734
});
2745-
this.d.usageLimit.show(limitCause ? { cause: limitCause } : undefined);
2735+
this.d.usageLimit.show(
2736+
limitCause === "model_gate"
2737+
? {
2738+
cause: limitCause,
2739+
model:
2740+
extractGatedModel(errorMessage, errorDetails) ?? undefined,
2741+
}
2742+
: limitCause
2743+
? { cause: limitCause }
2744+
: undefined,
2745+
);
27462746
return { stopReason: "rate_limited" };
27472747
}
27482748

‎packages/ui/src/features/auth/useAuthSession.ts‎

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { useHostTRPCClient } from "@posthog/host-router/react";
2-
import { BILLING_FLAG, USAGE_BILLING_FLAG } from "@posthog/shared";
2+
import { BILLING_FLAG } from "@posthog/shared";
33
import { useSeatStore } from "@posthog/ui/features/billing/seatStore";
44
import { useFeatureFlag } from "@posthog/ui/features/feature-flags/useFeatureFlag";
55
import {
@@ -112,18 +112,16 @@ function useSeatSync(
112112
billingEnabled: boolean,
113113
): void {
114114
const queryClient = useQueryClient();
115-
// Seats are retired under usage-based billing — provisioning one would
116-
// error against the retired seat API.
117-
const usageBillingEnabled = useFeatureFlag(USAGE_BILLING_FLAG);
118115
useEffect(() => {
119-
if (!authIdentity || !billingEnabled || usageBillingEnabled) {
116+
if (!authIdentity || !billingEnabled) {
120117
useSeatStore.getState().reset();
121118
return;
122119
}
123120

121+
// No-ops (and resets) once seats are retired — the store owns that gate.
124122
void useSeatStore.getState().fetchSeat({ autoProvision: true });
125123
void queryClient.invalidateQueries({ queryKey: [["llmGateway"]] });
126-
}, [authIdentity, billingEnabled, usageBillingEnabled, queryClient]);
124+
}, [authIdentity, billingEnabled, queryClient]);
127125
}
128126

129127
export function useAuthSession() {

0 commit comments

Comments
 (0)