Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions packages/core/src/billing/usageLimitContent.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import { describe, expect, it } from "vitest";
import { usageLimitContent } from "./usageLimitContent";

describe("usageLimitContent", () => {
it("offers a payment method for the model gate", () => {
const content = usageLimitContent({
cause: "model_gate",
resetLabel: null,
subscribed: false,
});
expect(content.title).toBe("Unlock premium models");
expect(content.description).toContain("This model isn't");
expect(content.actionLabel).toBe("Add payment method");
});

it.each([
// Confirmed-free org: allocation used up, the fix is adding a card.
[false, "Free usage used up", "Add payment method"],
// Subscribed org: the fix is raising the spend limit.
[true, "Organization usage limit reached", "Manage billing"],
// Unknown subscription state must not read as free.
[undefined, "Organization usage limit reached", "Manage billing"],
] as const)(
"org_limit with subscribed=%s -> %s / %s",
(subscribed, title, actionLabel) => {
const content = usageLimitContent({
cause: "org_limit",
resetLabel: null,
subscribed,
});
expect(content.title).toBe(title);
expect(content.actionLabel).toBe(actionLabel);
},
);

it("includes the reset hint in the free-tier copy when available", () => {
const content = usageLimitContent({
cause: "org_limit",
resetLabel: "Resets in 3h",
subscribed: false,
});
expect(content.title).toBe("Free usage used up");
expect(content.description).toContain("Resets in 3h");
});

it("renders generic copy without a billing CTA when the cause is unknown", () => {
const content = usageLimitContent({
cause: null,
resetLabel: "Resets in 2h",
subscribed: true,
});
expect(content.title).toBe("Usage limit reached");
expect(content.description).toContain("Resets in 2h");
expect(content.actionLabel).toBeNull();
expect(content.dismissLabel).toBe("Got it");
});
});
57 changes: 57 additions & 0 deletions packages/core/src/billing/usageLimitContent.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import type { GatewayLimitCause } from "@posthog/shared";

export interface UsageLimitContent {
title: string;
description: string;
actionLabel: string | null;
dismissLabel: string;
}

export function usageLimitContent(args: {
cause: GatewayLimitCause | null;
resetLabel: string | null;
subscribed: boolean | undefined;
}): UsageLimitContent {
const { cause, resetLabel, subscribed } = args;

if (cause === "model_gate") {
return {
title: "Unlock premium models",
description:
"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.",
actionLabel: "Add payment method",
dismissLabel: "Not now",
};
}

if (cause === "org_limit") {
if (subscribed === false) {
return {
title: "Free usage used up",
description: `Your organization has used up its included PostHog Code usage.${
resetLabel ? ` ${resetLabel}.` : ""
} Add a payment method to keep going — you only pay for what you use.`,
actionLabel: "Add payment method",
dismissLabel: "Not now",
};
}
return {
title: "Organization usage limit reached",
description:
"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.",
actionLabel: "Manage billing",
dismissLabel: "Got it",
};
}

// Not a billing denial (e.g. an upstream provider's own rate limit) —
// don't send the user to billing for something billing can't fix.
return {
title: "Usage limit reached",
description: `PostHog Code hit a usage limit.${
resetLabel ? ` ${resetLabel}.` : ""
} Please try again shortly.`,
actionLabel: null,
dismissLabel: "Got it",
};
}
187 changes: 185 additions & 2 deletions packages/core/src/llm-gateway/llm-gateway.test.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import { beforeEach, describe, expect, it, vi } from "vitest";
import type { AuthService } from "../auth/auth";
import type {
LlmGatewayAuth,
LlmGatewayEndpoints,
Expand Down Expand Up @@ -43,8 +44,25 @@ function createService(
};
const logger = { ...log, scope: () => log };

const service = new LlmGatewayService(host, logger);
return { service, auth, endpoints, log };
const orgListeners: Array<(state: { currentOrgId: string | null }) => void> =
[];
const authService = {
getState: () => ({ currentOrgId: "org-1" }),
on: (
_event: string,
listener: (state: { currentOrgId: string | null }) => void,
) => {
orgListeners.push(listener);
},
} as unknown as AuthService;
const emitAuthState = (currentOrgId: string | null) => {
for (const listener of orgListeners) {
listener({ currentOrgId });
}
};

const service = new LlmGatewayService(host, logger, authService);
return { service, auth, endpoints, log, emitAuthState };
}

const SUCCESS_BODY = {
Expand Down Expand Up @@ -157,6 +175,137 @@ describe("LlmGatewayService.prompt", () => {
});
});

it("surfaces a FastAPI bare-string detail as the error message", async () => {
const fetchMock = vi.fn().mockResolvedValue(
createJsonResponse(
{
detail: "OAuth application not authorized for product 'posthog_code'",
},
403,
),
);
const { service } = createService(fetchMock);

await expect(
service.prompt([{ role: "user", content: "hi" }]),
).rejects.toMatchObject({
name: "LlmGatewayError",
message: "OAuth application not authorized for product 'posthog_code'",
type: "unknown_error",
statusCode: 403,
});
});

// The free-tier model gate's 403 body, as the gateway serves it.
const MODEL_GATE_BODY = {
error: {
message:
"Model 'claude-haiku-4-5' needs a paid PostHog plan. Models available on the free tier: @cf/zai-org/glm-5.2. Add a payment method to your organization to unlock all models. (rate_limit)",
type: "permission_error",
code: "model_gate",
},
};

it("retries once on the free-tier model when the model gate 403s", async () => {
const fetchMock = vi
.fn()
.mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
.mockResolvedValueOnce(createJsonResponse(SUCCESS_BODY));
const { service } = createService(fetchMock);

const result = await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});

expect(result.content).toBe("hello world");
expect(fetchMock).toHaveBeenCalledTimes(2);
const firstBody = JSON.parse(fetchMock.mock.calls[0][1].body as string);
const retryBody = JSON.parse(fetchMock.mock.calls[1][1].body as string);
expect(firstBody.model).toBe("claude-haiku-4-5");
expect(retryBody.model).toBe("@cf/zai-org/glm-5.2");
});

it("routes straight to the free-tier model once the org is known unsubscribed", async () => {
const fetchMock = vi
.fn()
.mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
.mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
const { service } = createService(fetchMock);

// First call learns "unsubscribed" from the gate's 403.
await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});
// Second call must not burn a round trip on the gate.
await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});

expect(fetchMock).toHaveBeenCalledTimes(3);
const thirdBody = JSON.parse(fetchMock.mock.calls[2][1].body as string);
expect(thirdBody.model).toBe("@cf/zai-org/glm-5.2");
});

it("forgets the learned subscription state when the organization changes", async () => {
const fetchMock = vi
.fn()
.mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
.mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
const { service, emitAuthState } = createService(fetchMock);

await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});
emitAuthState("org-2");
await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});

const bodyAfterSwitch = JSON.parse(
fetchMock.mock.calls[2][1].body as string,
);
expect(bodyAfterSwitch.model).toBe("claude-haiku-4-5");
});

it("keeps the learned subscription state across same-org auth changes", async () => {
const fetchMock = vi
.fn()
.mockResolvedValueOnce(createJsonResponse(MODEL_GATE_BODY, 403))
.mockImplementation(async () => createJsonResponse(SUCCESS_BODY));
const { service, emitAuthState } = createService(fetchMock);

await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});
emitAuthState("org-1");
await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});

expect(fetchMock).toHaveBeenCalledTimes(3);
const bodyAfterRefresh = JSON.parse(
fetchMock.mock.calls[2][1].body as string,
);
expect(bodyAfterRefresh.model).toBe("@cf/zai-org/glm-5.2");
});

it("does not retry non-gate 403s on the free-tier model", async () => {
const fetchMock = vi
.fn()
.mockResolvedValue(
createJsonResponse(
{ detail: "Product 'posthog_code' requires OAuth authentication" },
403,
),
);
const { service } = createService(fetchMock);

await expect(
service.prompt([{ role: "user", content: "hi" }]),
).rejects.toMatchObject({ statusCode: 403 });
expect(fetchMock).toHaveBeenCalledTimes(1);
});

it("throws a timeout LlmGatewayError when the request aborts via the internal timeout", async () => {
const fetchMock = vi.fn((_url: string, init?: RequestInit) => {
return new Promise<Response>((_resolve, reject) => {
Expand Down Expand Up @@ -214,6 +363,40 @@ describe("LlmGatewayService.fetchUsage", () => {
statusCode: 503,
});
});

it("parses the usage-based billing fields when present", async () => {
const fetchMock = vi.fn().mockResolvedValue(
createJsonResponse({
...USAGE_BODY,
ai_credits: { exhausted: true },
code_usage_subscribed: true,
}),
);
const { service } = createService(fetchMock);

const usage = await service.fetchUsage();

expect(usage.ai_credits?.exhausted).toBe(true);
expect(usage.code_usage_subscribed).toBe(true);
});

it("feeds code_usage_subscribed into helper model routing", async () => {
const fetchMock = vi
.fn()
.mockResolvedValueOnce(
createJsonResponse({ ...USAGE_BODY, code_usage_subscribed: false }),
)
.mockResolvedValue(createJsonResponse(SUCCESS_BODY));
const { service } = createService(fetchMock);

await service.fetchUsage();
await service.prompt([{ role: "user", content: "hi" }], {
model: "claude-haiku-4-5",
});

const promptBody = JSON.parse(fetchMock.mock.calls[1][1].body as string);
expect(promptBody.model).toBe("@cf/zai-org/glm-5.2");
});
});

describe("LlmGatewayService.invalidatePlanCache", () => {
Expand Down
Loading
Loading