Skip to content

Commit 5fc7b89

Browse files
authored
fix: cap in-memory journal to prevent heap OOM under sustained load (#114)
## Summary `Journal.entries` was unbounded — every request across 26 handlers (~179 call sites: chat completions, messages, responses, gemini, bedrock, embeddings, images, speech, transcription, video, ollama, cohere, search, rerank, moderation, a2a, mcp, agui, ws-*, etc.) pushed a full `{ method, path, headers, body, response }` record and never evicted. At sustained prod traffic this grows heap ~3.8MB/sec → 4GB → OOM in ~18 minutes. Observed exactly that on `showcase-aimock` Railway service: deterministic 0→4GB heap growth then `FATAL ERROR: Reached heap limit Allocation failed`. Crash cascades to ~7 downstream showcase services that route through aimock via `OPENAI_BASE_URL`. ## Fix FIFO size cap on `Journal.entries` via new `JournalOptions.maxEntries`. Default: - Programmatic (`new Journal()` / `createServer()`): **unbounded** (backwards-compat; 100+ test/library callers depend on this). - CLI: **1000** (default for `serve` / the GHCR image). Override via `--journal-max <N>`; `0` or omitted = unbounded. Eviction is a single `shift()` per over-cap add. At cap=1000 × ~5KB/entry ≈ 5MB steady-state — well under heap limits. ## Test plan - [x] 6 new red-green tests in `src/__tests__/journal.test.ts` (cap behavior, FIFO ordering, uncapped default, `getLast`/`findByFixture` post-eviction, 100k-add cap invariant) - [x] Full suite: 2449 tests pass - [x] Lint + build + prettier clean - [ ] Post-merge: GHCR image republishes with the fix; Railway picks up on next drift-rebuild or manual redeploy ## Follow-ups (separate PR) - CLI: reject negative `--journal-max` values (currently silently treated as unbounded) - `createServer()` default: flip to finite cap so long-running embedders don't inherit the leak - `fixtureMatchCountsByTestId` Map cap (narrower but also unbounded) - Correct the "amortized O(1)" comment — `Array.shift` is O(n); true O(1) would need a ring buffer
2 parents 1c1d979 + 8455e55 commit 5fc7b89

8 files changed

Lines changed: 122 additions & 48 deletions

File tree

CHANGELOG.md

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,11 @@
11
# @copilotkit/aimock
22

3+
## 1.14.1
4+
5+
### Patch Changes
6+
7+
- Cap in-memory journal (and fixture-match-counts map) to prevent heap OOM under sustained load. `Journal.entries` was unbounded, causing heap growth ~3.8MB/sec to 4GB → OOM in ~18 minutes on production Railway deployments. Default cap for CLI (`serve`) is now 1000 entries; programmatic `createServer()` remains unbounded by default (back-compat). See `--journal-max` flag.
8+
39
## 1.14.0
410

511
### Minor Changes

package-lock.json

Lines changed: 2 additions & 46 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "@copilotkit/aimock",
3-
"version": "1.14.0",
3+
"version": "1.14.1",
44
"description": "Mock infrastructure for AI application testing — LLM APIs, image generation, text-to-speech, transcription, video generation, MCP tools, A2A agents, AG-UI event streams, vector databases, search, rerank, and moderation. One package, one port, zero dependencies.",
55
"license": "MIT",
66
"keywords": [

src/__tests__/journal.test.ts

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,4 +308,69 @@ describe("Journal", () => {
308308
expect(journal.getLast()!.id).toBe(entry.id);
309309
});
310310
});
311+
312+
describe("maxEntries (FIFO eviction)", () => {
313+
it("caps entries to maxEntries, dropping oldest (FIFO)", () => {
314+
const journal = new Journal({ maxEntries: 3 });
315+
316+
journal.add(makeEntry({ path: "/a" }));
317+
journal.add(makeEntry({ path: "/b" }));
318+
journal.add(makeEntry({ path: "/c" }));
319+
journal.add(makeEntry({ path: "/d" }));
320+
journal.add(makeEntry({ path: "/e" }));
321+
322+
expect(journal.size).toBe(3);
323+
const all = journal.getAll();
324+
expect(all.map((e) => e.path)).toEqual(["/c", "/d", "/e"]);
325+
});
326+
327+
it("does not cap when maxEntries is unset (backwards compat)", () => {
328+
const journal = new Journal();
329+
for (let i = 0; i < 5000; i++) journal.add(makeEntry({ path: `/${i}` }));
330+
expect(journal.size).toBe(5000);
331+
});
332+
333+
it("treats maxEntries = 0 or negative as uncapped", () => {
334+
const journal0 = new Journal({ maxEntries: 0 });
335+
const journalNeg = new Journal({ maxEntries: -1 });
336+
for (let i = 0; i < 100; i++) {
337+
journal0.add(makeEntry({ path: `/${i}` }));
338+
journalNeg.add(makeEntry({ path: `/${i}` }));
339+
}
340+
expect(journal0.size).toBe(100);
341+
expect(journalNeg.size).toBe(100);
342+
});
343+
344+
it("getLast returns the most recent after eviction", () => {
345+
const journal = new Journal({ maxEntries: 2 });
346+
journal.add(makeEntry({ path: "/a" }));
347+
journal.add(makeEntry({ path: "/b" }));
348+
const last = journal.add(makeEntry({ path: "/c" }));
349+
expect(journal.getLast()!.id).toBe(last.id);
350+
expect(journal.getLast()!.path).toBe("/c");
351+
});
352+
353+
it("findByFixture only returns surviving entries after eviction", () => {
354+
const journal = new Journal({ maxEntries: 2 });
355+
const fixture: Fixture = { match: { userMessage: "x" }, response: { content: "X" } };
356+
357+
journal.add(makeEntry({ response: { status: 200, fixture } }));
358+
journal.add(makeEntry({ response: { status: 200, fixture } }));
359+
journal.add(makeEntry({ response: { status: 200, fixture } }));
360+
361+
expect(journal.findByFixture(fixture)).toHaveLength(2);
362+
});
363+
364+
it("memory does not grow unbounded under sustained load with cap", () => {
365+
// Red-green anchor for the leak fix: 100k adds with cap=500 must stay at 500.
366+
const journal = new Journal({ maxEntries: 500 });
367+
for (let i = 0; i < 100_000; i++) {
368+
journal.add(makeEntry({ path: `/${i}` }));
369+
}
370+
expect(journal.size).toBe(500);
371+
// Last 500 paths preserved, oldest 99,500 evicted
372+
expect(journal.getLast()!.path).toBe("/99999");
373+
expect(journal.getAll()[0].path).toBe("/99500");
374+
});
375+
});
311376
});

src/cli.ts

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ Options:
2525
--record Record mode: proxy unmatched requests and save fixtures
2626
--proxy-only Proxy mode: forward unmatched requests without saving
2727
--strict Strict mode: fail on unmatched requests
28+
--journal-max <n> Max request entries retained in memory (default: 1000, 0 = unbounded)
2829
--provider-openai <url> Upstream URL for OpenAI (used with --record)
2930
--provider-anthropic <url> Upstream URL for Anthropic
3031
--provider-gemini <url> Upstream URL for Gemini
@@ -70,6 +71,7 @@ const { values } = parseArgs({
7071
"chaos-drop": { type: "string" },
7172
"chaos-malformed": { type: "string" },
7273
"chaos-disconnect": { type: "string" },
74+
"journal-max": { type: "string", default: "1000" },
7375
help: { type: "boolean", default: false },
7476
},
7577
strict: true,
@@ -110,6 +112,12 @@ if (Number.isNaN(chunkSize) || chunkSize < 1) {
110112
process.exit(1);
111113
}
112114

115+
const journalMax = Number(values["journal-max"]);
116+
if (Number.isNaN(journalMax) || !Number.isInteger(journalMax)) {
117+
console.error(`Invalid journal-max: ${values["journal-max"]} (must be an integer)`);
118+
process.exit(1);
119+
}
120+
113121
const logger = new Logger(logLevel);
114122

115123
// Parse chaos config from CLI flags
@@ -256,6 +264,7 @@ async function main() {
256264
metrics: values.metrics,
257265
record,
258266
strict: values.strict,
267+
journalMaxEntries: journalMax,
259268
},
260269
mounts,
261270
);

src/journal.ts

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,31 @@ function matchCriteriaEqual(a: FixtureMatch, b: FixtureMatch): boolean {
2929
);
3030
}
3131

32+
export interface JournalOptions {
33+
/**
34+
* Maximum number of entries to retain. When exceeded, oldest entries are
35+
* dropped FIFO. Set to 0 (or a negative value) for unbounded retention
36+
* (the historical default — suitable for short-lived test runs only).
37+
*
38+
* Long-running servers (e.g. mock proxies in CI/demo environments) should
39+
* always set a finite cap: every request appends an entry holding the
40+
* request body + headers + fixture reference, and without a cap the
41+
* journal grows until the process OOMs.
42+
*/
43+
maxEntries?: number;
44+
}
45+
3246
export class Journal {
3347
private entries: JournalEntry[] = [];
3448
private readonly fixtureMatchCountsByTestId: Map<string, Map<Fixture, number>> = new Map();
49+
private readonly maxEntries: number;
50+
51+
constructor(options: JournalOptions = {}) {
52+
// Treat 0 or negative as "unbounded" to preserve prior behavior when
53+
// the option is omitted or explicitly disabled.
54+
const cap = options.maxEntries;
55+
this.maxEntries = cap !== undefined && cap > 0 ? cap : 0;
56+
}
3557

3658
/** Backwards-compatible accessor — returns the default (no testId) count map. */
3759
get fixtureMatchCounts(): Map<Fixture, number> {
@@ -45,6 +67,12 @@ export class Journal {
4567
...entry,
4668
};
4769
this.entries.push(full);
70+
// FIFO eviction when over capacity. shift() in a tight loop would be
71+
// O(n^2); we only ever overshoot by one per add, so a single shift is
72+
// amortized O(1) per request.
73+
if (this.maxEntries > 0 && this.entries.length > this.maxEntries) {
74+
this.entries.shift();
75+
}
4876
return full;
4977
}
5078

src/server.ts

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -715,7 +715,7 @@ export async function createServer(
715715
}
716716
}
717717

718-
const journal = new Journal();
718+
const journal = new Journal({ maxEntries: options?.journalMaxEntries });
719719
const videoStates: VideoStateMap = new Map();
720720

721721
// Share journal and metrics registry with mounted services

src/types.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -398,6 +398,16 @@ export interface MockServerOptions {
398398
strict?: boolean;
399399
/** Record-and-replay: proxy unmatched requests to upstream and save fixtures. */
400400
record?: RecordConfig;
401+
/**
402+
* Maximum number of request/response entries to retain in the in-memory
403+
* journal. Oldest entries are dropped FIFO when the cap is exceeded.
404+
* Set to 0 (or a negative value) for unbounded retention.
405+
*
406+
* Defaults vary by invocation path: the CLI applies a finite cap suitable
407+
* for long-running servers (see `cli.ts`); programmatic callers default
408+
* to unbounded to preserve prior behavior for short-lived test runs.
409+
*/
410+
journalMaxEntries?: number;
401411
/**
402412
* Normalize requests before matching and recording. Useful for stripping
403413
* dynamic data (timestamps, UUIDs, session IDs) that would cause fixture

0 commit comments

Comments
 (0)