Skip to content

Commit 75c439f

Browse files
authored
fix(usage): bound OpenCode history reads to prevent backend OOM (#17961)
1 parent c727d12 commit 75c439f

2 files changed

Lines changed: 112 additions & 19 deletions

File tree

‎packages/provider-opencode/src/server/usage.test.ts‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import * as FileSystem from "effect/FileSystem";
77
import * as Layer from "effect/Layer";
88
import * as Path from "effect/Path";
99
import * as SqlClient from "effect/sql/SqlClient";
10+
import { vi } from "vite-plus/test";
1011

1112
import { readOpenCodeUsage } from "./usage.ts";
1213

@@ -17,6 +18,75 @@ const openDatabase = (filename: string) =>
1718
);
1819

1920
describe("readOpenCodeUsage", () => {
21+
it.effect("bounds raw history allocations while counting every eligible message once", () =>
22+
Effect.gen(function* () {
23+
const fileSystem = yield* FileSystem.FileSystem;
24+
const path = yield* Path.Path;
25+
const dir = yield* fileSystem.makeTempDirectoryScoped({ prefix: "usage-reader-large-test-" });
26+
const db = yield* openDatabase(path.join(dir, "opencode.db"));
27+
yield* db.unsafe(
28+
"CREATE TABLE message (id TEXT, session_id TEXT, data TEXT, time_created INTEGER)",
29+
);
30+
const sinceMs = 1780000000000;
31+
const message = {
32+
role: "assistant",
33+
modelID: "claude-sonnet-4-5",
34+
time: { created: sinceMs },
35+
tokens: { input: 100, output: 20 },
36+
// Transcript bodies are irrelevant to usage, but can be very large.
37+
body: "x".repeat(64 * 1024),
38+
};
39+
const data = JSON.stringify(message);
40+
const insert = (rowId: number, id: string, payload: string, created = sinceMs) =>
41+
db.unsafe(
42+
"INSERT INTO message (rowid, id, session_id, data, time_created) VALUES (?, ?, ?, ?, ?)",
43+
[rowId, id, "session-1", payload, created],
44+
);
45+
// Sparse IDs, including negative rowids, must not skip records at page boundaries.
46+
for (let index = 0; index < 130; index++) {
47+
yield* insert(index * 2 - 10, `msg-${index}`, data);
48+
}
49+
yield* insert(300, "old", data, sinceMs - 1);
50+
yield* insert(301, "user", JSON.stringify({ ...message, role: "user" }));
51+
yield* insert(302, "invalid", "{invalid");
52+
yield* insert(303, "msg-0", data);
53+
54+
let largestRawResult = 0;
55+
const originalAll = NodeSqlite.StatementSync.prototype.all;
56+
yield* Effect.acquireRelease(
57+
Effect.sync(() =>
58+
vi.spyOn(NodeSqlite.StatementSync.prototype, "all").mockImplementation(function (
59+
this: NodeSqlite.StatementSync,
60+
...params
61+
) {
62+
const rows = originalAll.apply(this, params);
63+
const rawSize = rows.reduce(
64+
(size, row) => size + (typeof row.data === "string" ? row.data.length : 0),
65+
0,
66+
);
67+
largestRawResult = Math.max(largestRawResult, rawSize);
68+
return rows;
69+
}),
70+
),
71+
(spy) => Effect.sync(() => spy.mockRestore()),
72+
);
73+
const result = yield* readOpenCodeUsage(dir, sinceMs);
74+
assert.isFalse(result.error);
75+
const records = result.files.flatMap((file) => file.records);
76+
assert.deepStrictEqual(
77+
records.map((record) => record.dedupeKey),
78+
Array.from({ length: 130 }, (_, index) => `opencode:msg-${index}`),
79+
);
80+
assert.strictEqual(
81+
records.reduce((total, record) => total + record.totals.uncachedInputTokens, 0),
82+
13000,
83+
);
84+
// The fixture exceeds 8 MiB. No query may retain its complete bodies.
85+
assert.isAbove(largestRawResult, 0);
86+
assert.isBelow(largestRawResult, 6 * 1024 * 1024);
87+
}).pipe(Effect.provide(NodeServices.layer)),
88+
);
89+
2090
it.effect("counts migrated OpenCode messages once and sees subsequent WAL writes", () =>
2191
Effect.gen(function* () {
2292
const fileSystem = yield* FileSystem.FileSystem;
@@ -68,3 +138,4 @@ describe("readOpenCodeUsage", () => {
68138
}).pipe(Effect.provide(NodeServices.layer)),
69139
);
70140
});
141+
import * as NodeSqlite from "node:sqlite";

‎packages/provider-opencode/src/server/usage.ts‎

Lines changed: 41 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -91,6 +91,7 @@ export interface OpenCodeUsageReadResult {
9191
}
9292

9393
const isNotFound = (cause: PlatformError.PlatformError) => cause.reason._tag === "NotFound";
94+
const MESSAGE_PAGE_SIZE = 64;
9495

9596
/** A regular file or directory, never a symlink to one. */
9697
const entryType = Effect.fn("entryType")(function* (path: string) {
@@ -172,26 +173,47 @@ export const readOpenCodeUsage = Effect.fn("readOpenCodeUsage")(function* (
172173
const timestamp = columns.has("time_created") ? "time_created" : "NULL";
173174
const predicates = table === "session_message" ? ["type = 'assistant'"] : [];
174175
if (timestamp !== "NULL") predicates.push("time_created >= ?");
175-
const where = predicates.length > 0 ? ` WHERE ${predicates.join(" AND ")}` : "";
176-
const rows = yield* sql.unsafe<{
177-
readonly id: unknown;
178-
readonly session_id: unknown;
179-
readonly data: unknown;
180-
readonly created: unknown;
181-
}>(
182-
`SELECT id, session_id, data, ${timestamp} AS created FROM ${table}${where}`,
183-
timestamp === "NULL" ? [] : [sinceMs],
184-
);
185-
for (const [index, row] of rows.entries()) {
186-
append(
187-
file.records,
188-
parseOpenCodeMessage(text(row.data), {
189-
id: text(row.id),
190-
sessionId: text(row.session_id),
191-
...(typeof row.created === "number" ? { timestampMs: row.created } : {}),
192-
}),
176+
const lastRowId = (yield* sql.unsafe<{ readonly last_row_id: number | null }>(
177+
`SELECT MAX(rowid) AS last_row_id FROM ${table}`,
178+
))[0]?.last_row_id;
179+
if (lastRowId === null || lastRowId === undefined) continue;
180+
let afterRowId: number | undefined;
181+
// SqlClient materializes each result with StatementSync.all(). Keep
182+
// message bodies in small pages, and exclude rows appended mid-scan.
183+
while (true) {
184+
const pagePredicates = [
185+
...predicates,
186+
"rowid <= ?",
187+
...(afterRowId === undefined ? [] : ["rowid > ?"]),
188+
];
189+
const rows = yield* sql.unsafe<{
190+
readonly row_id: number;
191+
readonly id: unknown;
192+
readonly session_id: unknown;
193+
readonly data: unknown;
194+
readonly created: unknown;
195+
}>(
196+
`SELECT rowid AS row_id, id, session_id, data, ${timestamp} AS created
197+
FROM ${table} WHERE ${pagePredicates.join(" AND ")} ORDER BY rowid LIMIT ${MESSAGE_PAGE_SIZE}`,
198+
[
199+
...(timestamp === "NULL" ? [] : [sinceMs]),
200+
lastRowId,
201+
...(afterRowId === undefined ? [] : [afterRowId]),
202+
],
193203
);
194-
if (index % 256 === 255) yield* Effect.yieldNow;
204+
for (const row of rows) {
205+
append(
206+
file.records,
207+
parseOpenCodeMessage(text(row.data), {
208+
id: text(row.id),
209+
sessionId: text(row.session_id),
210+
...(typeof row.created === "number" ? { timestampMs: row.created } : {}),
211+
}),
212+
);
213+
}
214+
if (rows.length < MESSAGE_PAGE_SIZE) break;
215+
afterRowId = rows[rows.length - 1]!.row_id;
216+
yield* Effect.yieldNow;
195217
}
196218
}
197219
}).pipe(

0 commit comments

Comments
 (0)