Skip to content

Commit 3121c19

Browse files
maria-rcksTrench-Monster
maria-rcks
and
Trench-Monster
authored
fix(usage): count compressed Codex rollouts (#25453)
Co-authored-by: Michael <[email protected]>
1 parent f048a9b commit 3121c19

3 files changed

Lines changed: 187 additions & 7 deletions

File tree

‎apps/server/src/usage/UsageService.test.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ import * as NodeFSP from "node:fs/promises";
55
import * as NodeOS from "node:os";
66
import * as NodePath from "node:path";
77
import * as NodeSqlite from "node:sqlite";
8+
import * as NodeZlib from "node:zlib";
89

910
import { assert, describe, it } from "@effect/vitest";
1011
import { vi } from "vite-plus/test";
@@ -467,6 +468,73 @@ describe("Codex usage history", () => {
467468
});
468469

469470
describe("UsageService", () => {
471+
it.live("counts compressed Codex rollouts across refresh, rewrites, and cache restarts", () =>
472+
Effect.gen(function* () {
473+
const { home, settings } = yield* setup;
474+
const sessions = NodePath.join(home, "codex", "sessions");
475+
const compressedPath = NodePath.join(sessions, "compressed.jsonl.zst");
476+
const usage = (output_tokens: number, timestamp = "2026-08-01T10:00:00Z") =>
477+
encodeUnknownJsonString({
478+
type: "event_msg",
479+
timestamp,
480+
payload: {
481+
type: "token_count",
482+
info: { last_token_usage: { input_tokens: 10, output_tokens } },
483+
},
484+
}) + "\n";
485+
const rollout = (id: string, output: number) =>
486+
[
487+
{ type: "session_meta", payload: { id } },
488+
{ type: "turn_context", payload: { model: "gpt-6-astra" } },
489+
]
490+
.map((line) => encodeUnknownJsonString(line))
491+
.join("\n") +
492+
"\n" +
493+
usage(output);
494+
const compressed = rollout("compressed-session", 23);
495+
yield* Effect.promise(async () => {
496+
await NodeFSP.mkdir(sessions, { recursive: true });
497+
await NodeFSP.writeFile(
498+
NodePath.join(sessions, "plain.jsonl"),
499+
rollout("plain-session", 17),
500+
);
501+
await NodeFSP.writeFile(compressedPath, NodeZlib.zstdCompressSync(compressed));
502+
});
503+
yield* Effect.gen(function* () {
504+
const service = yield* UsageService.make;
505+
const first = yield* service.readSummary(WINDOW);
506+
assert.strictEqual(totalOutputTokens(first), 40);
507+
assert.strictEqual(
508+
first.sources.find((source) => source.fingerprint.provider === "codex")?.distinctSessions,
509+
2,
510+
);
511+
assert.deepStrictEqual((yield* service.readSummary(WINDOW)).buckets, first.buckets);
512+
yield* service.awaitPersisted;
513+
// The same rollout may remain in a plain file while compression finishes.
514+
yield* Effect.promise(() =>
515+
NodeFSP.writeFile(NodePath.join(sessions, "copy.jsonl"), compressed),
516+
);
517+
assert.strictEqual(totalOutputTokens(yield* service.readSummary(WINDOW)), 40);
518+
yield* Effect.promise(() =>
519+
NodeFSP.writeFile(
520+
compressedPath,
521+
NodeZlib.zstdCompressSync(compressed + usage(31, "2026-08-01T10:00:05Z")),
522+
),
523+
);
524+
const rewritten = yield* service.readSummary(WINDOW);
525+
assert.strictEqual(totalOutputTokens(rewritten), 71);
526+
assert.deepStrictEqual((yield* service.readSummary(WINDOW)).buckets, rewritten.buckets);
527+
yield* service.awaitPersisted;
528+
const restarted = yield* UsageService.make;
529+
assert.deepStrictEqual((yield* restarted.readSummary(WINDOW)).buckets, rewritten.buckets);
530+
yield* restarted.awaitPersisted;
531+
}).pipe(
532+
Effect.scoped,
533+
Effect.provide(layerService({ prefix: "usage-compressed-codex", home, settings })),
534+
);
535+
}).pipe(Effect.scoped),
536+
);
537+
470538
it.live("keeps the first host identity across recreation and explicit override removal", () =>
471539
Effect.gen(function* () {
472540
const { transcript, settings, home } = yield* setup;

‎apps/server/src/usage/usageTranscriptReader.test.ts‎

Lines changed: 103 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,18 @@
44
import * as NodeFSP from "node:fs/promises";
55
import * as NodeOS from "node:os";
66
import * as NodePath from "node:path";
7+
import * as NodeZlib from "node:zlib";
78

89
import { afterEach, assert, beforeEach, describe, it } from "@effect/vitest";
910
import type { TranscriptUsageFormat } from "@t3tools/provider-core/server/usage";
1011
import * as Schema from "effect/Schema";
1112

1213
import { TEST_FORMATS } from "./usageTestFormats.ts";
13-
import { readTranscriptRecords, transcriptLineGate } from "./usageTranscriptReader.ts";
14+
import {
15+
listTranscriptFiles,
16+
readTranscriptRecords,
17+
transcriptLineGate,
18+
} from "./usageTranscriptReader.ts";
1419

1520
let dir: string;
1621

@@ -71,6 +76,103 @@ function codexUsageLine(outputTokens: number, secondsOffset: number): string {
7176
})}\n`;
7277
}
7378

79+
describe("compressed transcripts", () => {
80+
it("discovers nested plain and zstd JSONL files within the requested window", async () => {
81+
const nested = NodePath.join(dir, "nested");
82+
await NodeFSP.mkdir(nested);
83+
const plain = NodePath.join(dir, "active.jsonl");
84+
const compressed = NodePath.join(nested, "archived.jsonl.zst");
85+
const old = NodePath.join(nested, "old.jsonl.zst");
86+
await NodeFSP.writeFile(plain, claudeLine(1, 5));
87+
await NodeFSP.writeFile(compressed, NodeZlib.zstdCompressSync(claudeLine(2, 7)));
88+
await NodeFSP.writeFile(old, NodeZlib.zstdCompressSync(claudeLine(3, 11)));
89+
await NodeFSP.writeFile(NodePath.join(nested, "unrelated.zst"), "ignored");
90+
await NodeFSP.writeFile(NodePath.join(nested, "partial.jsonl.zst.tmp"), "ignored");
91+
await NodeFSP.utimes(old, 1, 1);
92+
93+
const listed = await listTranscriptFiles(dir, Date.parse("2026-01-01T00:00:00Z"));
94+
assert.deepStrictEqual(
95+
listed.files.map((file) => file.path).sort(),
96+
[plain, compressed].sort(),
97+
);
98+
assert.strictEqual(listed.failedPaths, 0);
99+
assert.isAbove(listed.files.find((file) => file.path === compressed)?.size ?? 0, 0);
100+
});
101+
102+
it.each([
103+
{ parser: "native", streamingThresholdBytes: undefined },
104+
{ parser: "projected", streamingThresholdBytes: 1024 },
105+
])(
106+
"matches plain Codex records and unfinished tails through long lines ($parser)",
107+
async ({ streamingThresholdBytes }) => {
108+
const plain = NodePath.join(dir, "rollout.jsonl");
109+
const compressed = plain + ".zst";
110+
const large = codexUsageLine(9, 5).replace(
111+
'"payload":',
112+
'"padding":' + JSON.stringify("日本語 é".repeat(128 * 1024)) + ',"payload":',
113+
);
114+
const content =
115+
codexMetaLine() +
116+
codexModelLine("gpt-6-astra") +
117+
codexTierLine("ultrafast") +
118+
large +
119+
codexUsageLine(13, 6).replace('"output_tokens":13', '"output_tokens":"invalid"') +
120+
codexUsageLine(21, 8).trimEnd();
121+
await NodeFSP.writeFile(plain, content);
122+
await NodeFSP.writeFile(compressed, NodeZlib.zstdCompressSync(content));
123+
124+
const options =
125+
streamingThresholdBytes === undefined ? undefined : { streamingThresholdBytes };
126+
const expected = await readTranscriptRecords(plain, TEST_FORMATS.codex, undefined, options);
127+
const actual = await readTranscriptRecords(
128+
compressed,
129+
TEST_FORMATS.codex,
130+
undefined,
131+
options,
132+
);
133+
assert.isNotNull(expected);
134+
assert.isNotNull(actual);
135+
assert.strictEqual(expected.records.length, 1);
136+
assert.strictEqual(expected.tailRecords.length, 1);
137+
assert.deepStrictEqual(actual.records, expected.records);
138+
assert.deepStrictEqual(actual.tailRecords, expected.tailRecords);
139+
assert.strictEqual(expected.malformedRecords, 1);
140+
assert.strictEqual(actual.malformedRecords, expected.malformedRecords);
141+
assert.strictEqual(actual.records[0]?.speed, "ultrafast");
142+
assert.isFalse(actual.resumed);
143+
assert.strictEqual(actual.position.resumeOffset, 0);
144+
},
145+
);
146+
147+
it("ignores a prior resume position and rereads compressed rollouts from their start", async () => {
148+
const plain = NodePath.join(dir, "rollout.jsonl");
149+
const content = codexMetaLine() + codexModelLine("gpt-6-astra") + codexUsageLine(9, 5);
150+
await NodeFSP.writeFile(plain, content);
151+
const first = await readTranscriptRecords(plain, TEST_FORMATS.codex);
152+
assert.isNotNull(first);
153+
154+
const compressed = plain + ".zst";
155+
await NodeFSP.writeFile(compressed, NodeZlib.zstdCompressSync(content + codexUsageLine(21, 8)));
156+
const actual = await readTranscriptRecords(compressed, TEST_FORMATS.codex, first.position);
157+
assert.isNotNull(actual);
158+
assert.isFalse(actual.resumed);
159+
assert.deepStrictEqual(
160+
actual.records.map((record) => record.totals.outputTokens),
161+
[9, 21],
162+
);
163+
const repeated = await readTranscriptRecords(compressed, TEST_FORMATS.codex, actual.position);
164+
assert.isNotNull(repeated);
165+
assert.isFalse(repeated.resumed);
166+
assert.deepStrictEqual(repeated.records, actual.records);
167+
});
168+
169+
it("returns null for zstd decoder errors rather than caching an empty parse", async () => {
170+
const path = NodePath.join(dir, "corrupt.jsonl.zst");
171+
await NodeFSP.writeFile(path, "not a zstd frame");
172+
assert.isNull(await readTranscriptRecords(path, TEST_FORMATS.claude));
173+
});
174+
});
175+
74176
describe("readTranscriptRecords resume", () => {
75177
it.each([
76178
{ mode: "buffered", streamingThresholdBytes: Infinity },

‎apps/server/src/usage/usageTranscriptReader.ts‎

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import * as NodeFSP from "node:fs/promises";
44
import * as NodePath from "node:path";
55
import * as NodeStringDecoder from "node:string_decoder";
6+
import * as NodeZlib from "node:zlib";
67

78
import type {
89
SelectedFields,
@@ -85,7 +86,11 @@ export async function listTranscriptFiles(
8586
for (const entry of entries) {
8687
const child = NodePath.join(dir, entry.name);
8788
if (entry.isDirectory()) await walk(child);
88-
else if (fileName !== undefined ? entry.name === fileName : entry.name.endsWith(".jsonl")) {
89+
else if (
90+
fileName !== undefined
91+
? entry.name === fileName
92+
: entry.name.endsWith(".jsonl") || entry.name.endsWith(".jsonl.zst")
93+
) {
8994
candidates.push(child);
9095
}
9196
}
@@ -162,6 +167,7 @@ export async function readTranscriptRecords<State>(
162167
options?: { readonly streamingThresholdBytes?: number },
163168
): Promise<TranscriptParseResult | null> {
164169
const streamingThresholdBytes = options?.streamingThresholdBytes ?? STREAMING_THRESHOLD_BYTES;
170+
const compressed = filePath.endsWith(".jsonl.zst");
165171
let handle: NodeFSP.FileHandle;
166172
try {
167173
handle = await NodeFSP.open(filePath, "r");
@@ -181,6 +187,7 @@ export async function readTranscriptRecords<State>(
181187
? Option.none()
182188
: decodeState(resumeFrom.state);
183189
if (
190+
!compressed &&
184191
resumeFrom !== undefined &&
185192
resumeFrom.resumeOffset > 0 &&
186193
(codec === undefined || Option.isSome(savedState)) &&
@@ -247,11 +254,14 @@ export async function readTranscriptRecords<State>(
247254
streaming = undefined;
248255
decoder = undefined;
249256
};
250-
const stream = handle.createReadStream({
257+
const input = handle.createReadStream({
251258
start,
252-
autoClose: false,
259+
autoClose: compressed,
253260
highWaterMark: 256 * 1024,
254-
}) as AsyncIterable<Buffer>;
261+
});
262+
// Compressed snapshots cache a full parse; decoded byte offsets cannot
263+
// resume the compressed input.
264+
const stream = compressed ? input.compose(NodeZlib.createZstdDecompress()) : input;
255265
for await (const chunk of stream) {
256266
let lineStart = 0;
257267
while (lineStart < chunk.length) {
@@ -279,7 +289,7 @@ export async function readTranscriptRecords<State>(
279289
decodeState === undefined ? state : Option.getOrThrow(decodeState(encodedState));
280290
finish(tailState, tailRecords);
281291

282-
const guardLength = Math.min(GUARD_LENGTH, resumeOffset);
292+
const guardLength = compressed ? 0 : Math.min(GUARD_LENGTH, resumeOffset);
283293
let guardHash = 0;
284294
if (guardLength > 0) {
285295
const window = Buffer.alloc(guardLength);
@@ -292,7 +302,7 @@ export async function readTranscriptRecords<State>(
292302
tailRecords,
293303
malformedRecords: format.readMalformedRecords?.(tailState) ?? 0,
294304
position: {
295-
resumeOffset,
305+
resumeOffset: compressed ? 0 : resumeOffset,
296306
guardLength,
297307
guardHash,
298308
state: encodedState,

0 commit comments

Comments
 (0)