Skip to content
55 changes: 55 additions & 0 deletions apps/server/src/provider/makeManagedServerProvider.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -657,4 +657,59 @@ describe("makeManagedServerProvider", () => {
}),
).pipe(Effect.provide(AlwaysRunTestLayer)),
);

it.effect("does not publish a window streamed after a failed probe as the full list", () =>
Effect.scoped(
Effect.gen(function* () {
const refreshCount = yield* Ref.make(0);
const failedLimits = {
checkedAt: "2026-04-10T00:00:01.000Z",
windows: [],
unavailable: { reason: "probeFailed" },
} as const;
const probedLimits = {
checkedAt: "2026-04-10T00:00:03.000Z",
windows: [
{ id: "five_hour", kind: "session", label: "Session", usedPercent: 10 },
{ id: "seven_day", kind: "weekly", label: "Weekly", usedPercent: 25 },
],
} as const;
const provider = yield* makeManagedServerProvider<TestSettings>({
resolveMaintenance: () => Effect.succeed(maintenanceCapabilities),
getSettings: Effect.succeed({ enabled: true }),
streamSettings: Stream.empty,
haveSettingsChanged: (previous, next) => previous.enabled !== next.enabled,
initialSnapshot: () => Effect.succeed(initialSnapshot),
checkProvider: Ref.updateAndGet(refreshCount, (count) => count + 1).pipe(
Effect.map((count) => ({
...refreshedSnapshot,
usageLimits: count < 3 ? failedLimits : probedLimits,
})),
),
refreshInterval: "1 hour",
});
yield* Stream.take(provider.streamChanges, 1).pipe(Stream.runDrain);

const weekly = {
id: "seven_day",
kind: "weekly",
label: "Weekly",
usedPercent: 20,
} as const;
yield* provider.applyUsageLimits({
checkedAt: "2026-04-10T00:00:02.000Z",
windows: [weekly],
});
assert.deepStrictEqual((yield* provider.getSnapshot).usageLimits, {
checkedAt: "2026-04-10T00:00:02.000Z",
windows: [weekly],
unavailable: { reason: "probeFailed" },
});

// A second failed probe has no last good list to keep.
assert.deepStrictEqual((yield* provider.refresh).usageLimits, failedLimits);
assert.deepStrictEqual((yield* provider.refresh).usageLimits, probedLimits);
}),
).pipe(Effect.provide(AlwaysRunTestLayer)),
);
});
28 changes: 28 additions & 0 deletions apps/server/src/provider/providerUsageLimits.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,34 @@ describe("applyUsageLimitsUpdate", () => {
).toBe(published);
});

it("keeps a failed probe marked when a sparse update lands on it", () => {
const failed = {
checkedAt,
windows: [],
unavailable: { reason: "probeFailed" as const, message: "Codex did not answer." },
};
const next = applyUsageLimitsUpdate({
previous: failed,
checkedAt: "2026-09-03T12:00:05.000Z",
update: { windows: [weekly] },
});
expect(next).toEqual({
checkedAt: "2026-09-03T12:00:05.000Z",
windows: [weekly],
unavailable: failed.unavailable,
});
// The same numbers again must not rebuild the snapshot and publish it.
expect(
applyUsageLimitsUpdate({ previous: next, checkedAt, update: { windows: [weekly] } }),
).toBe(next);
});

it("marks a sparse update with nothing published as a failed read", () => {
expect(
applyUsageLimitsUpdate({ previous: undefined, checkedAt, update: { windows: [weekly] } }),
).toEqual({ checkedAt, windows: [weekly], unavailable: { reason: "probeFailed" } });
});

it("preserves reset credits when a streamed window update changes usage", () => {
const resetCredits = { availableCount: 2, nextExpiresAt: "2026-10-01T00:00:00.000Z" };
const next = applyUsageLimitsUpdate({
Expand Down
10 changes: 9 additions & 1 deletion apps/server/src/provider/providerUsageLimits.ts
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,11 @@ export function makeUnavailableUsageLimits(input: {
*
* An `unsupported` snapshot stays unsupported: an account that cannot have
* subscription windows will not start reporting them mid-turn.
*
* An update never completes a read. Merged onto a failed probe, or onto no
* probe at all, it keeps the windows it names but leaves the snapshot
* `probeFailed`. Otherwise one streamed window would read as the account's
* full list, and the next failed probe would keep it as the last good one.
*/
export function applyUsageLimitsUpdate(input: {
readonly previous: ServerProviderUsageLimits | undefined;
Expand Down Expand Up @@ -88,12 +93,15 @@ export function applyUsageLimitsUpdate(input: {
changed = true;
}
}
if (!changed && previous !== undefined && previous.unavailable === undefined) {
if (!changed && previous !== undefined) {
return previous;
}
const unavailable =
previous === undefined ? { reason: "probeFailed" as const } : previous.unavailable;
return {
...makeUsageLimits({ checkedAt: input.checkedAt, windows: merged.values() }),
...(previous?.resetCredits !== undefined ? { resetCredits: previous.resetCredits } : {}),
...(unavailable !== undefined ? { unavailable } : {}),
};
}

Expand Down
Loading