Skip to content
Closed
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
2 changes: 2 additions & 0 deletions apps/code/src/main/di/container.ts
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@ import { SuspensionService } from "../services/suspension/service";
import { TaskLinkService } from "../services/task-link/service";
import { UIService } from "../services/ui/service";
import { UpdatesService } from "../services/updates/service";
import { UsageMonitorService } from "../services/usage-monitor/service";
import { WatcherRegistryService } from "../services/watcher-registry/service";
import { WorkspaceService } from "../services/workspace/service";
import { MAIN_TOKENS } from "./tokens";
Expand Down Expand Up @@ -146,6 +147,7 @@ container.bind(MAIN_TOKENS.ShellService).to(ShellService);
container.bind(MAIN_TOKENS.SlackIntegrationService).to(SlackIntegrationService);
container.bind(MAIN_TOKENS.UIService).to(UIService);
container.bind(MAIN_TOKENS.UpdatesService).to(UpdatesService);
container.bind(MAIN_TOKENS.UsageMonitorService).to(UsageMonitorService);
container.bind(MAIN_TOKENS.TaskLinkService).to(TaskLinkService);
container.bind(MAIN_TOKENS.InboxLinkService).to(InboxLinkService);
container.bind(MAIN_TOKENS.WatcherRegistryService).to(WatcherRegistryService);
Expand Down
1 change: 1 addition & 0 deletions apps/code/src/main/di/tokens.ts
Original file line number Diff line number Diff line change
Expand Up @@ -82,4 +82,5 @@ export const MAIN_TOKENS = Object.freeze({
ProvisioningService: Symbol.for("Main.ProvisioningService"),
WorkspaceService: Symbol.for("Main.WorkspaceService"),
EnrichmentService: Symbol.for("Main.EnrichmentService"),
UsageMonitorService: Symbol.for("Main.UsageMonitorService"),
});
28 changes: 28 additions & 0 deletions apps/code/src/main/services/usage-monitor/schemas.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
import { z } from "zod";

export const USAGE_THRESHOLDS = [50, 75, 90, 100] as const;
export type UsageThreshold = (typeof USAGE_THRESHOLDS)[number];

export const thresholdCrossedEvent = z.object({
bucket: z.enum(["burst", "sustained"]),
threshold: z.union([
z.literal(50),
z.literal(75),
z.literal(90),
z.literal(100),
]),
usedPercent: z.number(),
resetAt: z.string().datetime().nullable(),
resetsInSeconds: z.number(),
isPro: z.boolean(),
});

export type ThresholdCrossedEvent = z.infer<typeof thresholdCrossedEvent>;

export const UsageMonitorEvent = {
ThresholdCrossed: "threshold-crossed",
} as const;

export interface UsageMonitorEvents {
[UsageMonitorEvent.ThresholdCrossed]: ThresholdCrossedEvent;
}
182 changes: 182 additions & 0 deletions apps/code/src/main/services/usage-monitor/service.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,182 @@
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
import type { UsageOutput } from "../llm-gateway/schemas";
import { UsageMonitorEvent } from "./schemas";

const mockStoreGet = vi.hoisted(() => vi.fn());
const mockStoreSet = vi.hoisted(() => vi.fn());

vi.mock("./store", () => ({
usageMonitorStore: {
get: mockStoreGet,
set: mockStoreSet,
},
}));

vi.mock("../../utils/logger.js", () => ({
logger: {
scope: () => ({
info: vi.fn(),
error: vi.fn(),
warn: vi.fn(),
debug: vi.fn(),
}),
},
}));

import { LlmGatewayService } from "../llm-gateway/service";
import { UsageMonitorService } from "./service";

function makeUsage(overrides?: {
burstPercent?: number;
sustainedPercent?: number;
billingPeriodEnd?: string | null;
burstResetAt?: string;
sustainedResetAt?: string;
}): UsageOutput {
return {
product: "posthog_code",
user_id: 42,
is_rate_limited: false,
billing_period_end:
overrides?.billingPeriodEnd === undefined
? null
: overrides.billingPeriodEnd,
burst: {
used_percent: overrides?.burstPercent ?? 0,
resets_in_seconds: 3600,
reset_at: overrides?.burstResetAt ?? "2026-05-25T16:00:00.000Z",
exceeded: false,
},
sustained: {
used_percent: overrides?.sustainedPercent ?? 0,
resets_in_seconds: 86400,
reset_at: overrides?.sustainedResetAt ?? "2026-06-01T00:00:00.000Z",
exceeded: false,
},
};
}

function mockGateway(usage: UsageOutput | null): LlmGatewayService {
return {
fetchUsage: vi.fn().mockResolvedValue(usage),
} as unknown as LlmGatewayService;
}

describe("UsageMonitorService", () => {
let service: UsageMonitorService;
let persisted: Record<string, string>;

beforeEach(() => {
vi.useFakeTimers();
vi.setSystemTime(new Date("2026-05-25T12:00:00.000Z"));
persisted = {};
mockStoreGet.mockImplementation((_key: string, fallback: unknown) => ({
...persisted,
...(fallback as Record<string, string>),
}));
mockStoreSet.mockImplementation(
(_key: string, value: Record<string, string>) => {
persisted = { ...value };
},
);
});

afterEach(() => {
service?.stop();
vi.useRealTimers();
});

it("emits at 75% but not again on the next poll for the same anchor", async () => {
const events: unknown[] = [];
const gateway = mockGateway(makeUsage({ burstPercent: 78 }));
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));

await service.pollOnce();
expect(events).toHaveLength(1);
expect(events[0]).toMatchObject({
bucket: "burst",
threshold: 75,
usedPercent: 78,
});

await service.pollOnce();
expect(events).toHaveLength(1);
});

it("only emits the highest threshold a bucket has crossed", async () => {
const events: unknown[] = [];
const gateway = mockGateway(makeUsage({ burstPercent: 95 }));
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));

await service.pollOnce();
expect(events).toHaveLength(1);
expect(events[0]).toMatchObject({ threshold: 90 });
});

it("doesn't re-emit after a relaunch with persisted dedupe", async () => {
const events: unknown[] = [];
const gateway = mockGateway(makeUsage({ burstPercent: 55 }));
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));
await service.pollOnce();
expect(events).toHaveLength(1);
service.stop();

// Simulate relaunch
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));
await service.pollOnce();
expect(events).toHaveLength(1);
});

it("tracks burst and sustained as independent buckets", async () => {
const events: unknown[] = [];
const gateway = mockGateway(
makeUsage({
burstPercent: 55,
sustainedPercent: 80,
billingPeriodEnd: "2026-06-01T00:00:00.000Z",
}),
);
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));

await service.pollOnce();
expect(events).toHaveLength(2);
expect(events.map((e) => (e as { bucket: string }).bucket).sort()).toEqual([
"burst",
"sustained",
]);
});

it("marks events with isPro when billing_period_end is set", async () => {
const events: { isPro: boolean }[] = [];
const gateway = mockGateway(
makeUsage({
sustainedPercent: 60,
billingPeriodEnd: "2026-06-01T00:00:00.000Z",
}),
);
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) =>
events.push(e as { isPro: boolean }),
);

await service.pollOnce();
expect(events[0]?.isPro).toBe(true);
});

it("silently skips polls when the gateway throws", async () => {
const events: unknown[] = [];
const gateway = {
fetchUsage: vi.fn().mockRejectedValue(new Error("not authenticated")),
} as unknown as LlmGatewayService;
service = new UsageMonitorService(gateway);
service.on(UsageMonitorEvent.ThresholdCrossed, (e) => events.push(e));

await expect(service.pollOnce()).resolves.toBeNull();
expect(events).toHaveLength(0);
});
});
Comment on lines +89 to +182
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Missing escalation test; non-parameterised threshold cases

No test covers the scenario where usage crosses a lower threshold (e.g. 55% → 50% fires) and then rises to a higher one (85% → 75% fires) within the same anchor window. This is the most important in-window state-machine transition and it's currently untested.

Separately, the threshold detection cases ("emits at 75%" and "only emits the highest threshold") share the same structure and are good candidates for a parameterised test per the team's convention — e.g. it.each([[78, 75], [95, 90], [100, 100], [49, null]])("threshold at %i% fires %i", ...).

Prompt To Fix With AI
This is a comment left during a code review.
Path: apps/code/src/main/services/usage-monitor/service.test.ts
Line: 89-182

Comment:
**Missing escalation test; non-parameterised threshold cases**

No test covers the scenario where usage crosses a lower threshold (e.g. 55% → 50% fires) and then rises to a higher one (85% → 75% fires) within the same anchor window. This is the most important in-window state-machine transition and it's currently untested.

Separately, the threshold detection cases ("emits at 75%" and "only emits the highest threshold") share the same structure and are good candidates for a parameterised test per the team's convention — e.g. `it.each([[78, 75], [95, 90], [100, 100], [49, null]])("threshold at %i% fires %i", ...)`.

How can I resolve this? If you propose a fix, please make it concise.

Loading
Loading