fix(mcp): inject livez probe to kill mcp-standalone test flake (#449)

[rohitg00]

Closes #449.

DI shape chosen: option A (module-level singleton with override)

Added setLivezProbe(fn?: LivezProbe) in src/mcp/rest-proxy.ts. Tests
swap in an instant-fail probe before each call; default stays the real
fetch-based probe so production wire behaviour is identical.

Why option A over constructor DI (option B): the rest-proxy.ts module
is already singleton-shaped — resolveHandle() reads a module-scoped cache,
invalidateHandle() and resetHandleForTests() mutate that cache. The probe
is the same shape of dependency. A constructor seam would have meant
refactoring every caller of resolveHandle() (and the handleToolCall
chain in standalone.ts) to thread the proxy instance through. That's
beyond "DI seam, don't touch the logic." Option A keeps the diff to two
files and 84 added lines.

Also wired the override into resetHandleForTests() so a stub from one test
file can't leak into another.

10-run flake-kill validation

=== run 1 ===
      Tests  25 passed (25)
   Duration  185ms (transform 68ms, setup 0ms, import 87ms, tests 17ms, environment 0ms)
=== run 2 ===
      Tests  25 passed (25)
   Duration  188ms (transform 67ms, setup 0ms, import 85ms, tests 18ms, environment 0ms)
=== run 3 ===
      Tests  25 passed (25)
   Duration  190ms (transform 67ms, setup 0ms, import 85ms, tests 20ms, environment 0ms)
=== run 4 ===
      Tests  25 passed (25)
   Duration  190ms (transform 68ms, setup 0ms, import 87ms, tests 17ms, environment 0ms)
=== run 5 ===
      Tests  25 passed (25)
   Duration  188ms (transform 67ms, setup 0ms, import 85ms, tests 17ms, environment 0ms)
=== run 6 ===
      Tests  25 passed (25)
   Duration  180ms (transform 65ms, setup 0ms, import 82ms, tests 17ms, environment 0ms)
=== run 7 ===
      Tests  25 passed (25)
   Duration  192ms (transform 69ms, setup 0ms, import 87ms, tests 18ms, environment 0ms)
=== run 8 ===
      Tests  25 passed (25)
   Duration  182ms (transform 64ms, setup 0ms, import 81ms, tests 18ms, environment 0ms)
=== run 9 ===
      Tests  25 passed (25)
   Duration  185ms (transform 64ms, setup 0ms, import 82ms, tests 18ms, environment 0ms)
=== run 10 ===
      Tests  25 passed (25)
   Duration  183ms (transform 65ms, setup 0ms, import 83ms, tests 18ms, environment 0ms)

25 tests = 24 original + 1 new regression-guard (livez probe stub is invoked instead of the real fetch (issue #449)) that asserts the DI seam
is wired and the real network is never touched.

Pre-fix baseline was 2.2s per run with intermittent timeouts; post-fix is
~185ms with zero variance across 10 runs.

Full-suite stability

=== full-suite run 1 ===
 Test Files  90 passed (90)
      Tests  988 passed (988)
=== full-suite run 2 ===
 Test Files  90 passed (90)
      Tests  988 passed (988)
=== full-suite run 3 ===
 Test Files  90 passed (90)
      Tests  988 passed (988)

988 = the 987 main has today + 1 new regression-guard test in this PR.

mcp-standalone-proxy.test.ts (the sibling file that intentionally exercises
the proxy path via globalThis.fetch stubbing) still passes 14/14
unchanged — resetHandleForTests() restores the default probe in its
beforeEach/afterEach, so it never sees the stub installed by the
standalone file.

npm run build clean.

Diff stat

 src/mcp/rest-proxy.ts       | 44 ++++++++++++++++++++++++++++++++++++------
 test/mcp-standalone.test.ts | 47 ++++++++++++++++++++++++++++++++++++++++++++-
 2 files changed, 84 insertions(+), 7 deletions(-)

Constraints honoured

• No version bump
• No CHANGELOG edit
• No --no-verify
• No shim logic rewrite — probe behaviour (probe ok → proxy; probe fail →
  local KV) is byte-identical in production
• Branch fix/449-mcp-standalone-probe-di, conventional commit, not merged

Summary by CodeRabbit

• Tests
  • Enhanced test infrastructure with deterministic mocking for the liveness probe mechanism, preventing test contamination and improving reliability of proxy/fallback decision testing.

[vercel]

The latest updates on your projects. Learn more about Vercel for GitHub.

Project	Deployment	Actions	Updated (UTC)
agentmemory	Ready	Preview, Comment	May 17, 2026 10:52am

[coderabbitai]

📝 Walkthrough

Walkthrough

src/mcp/rest-proxy.ts exports a LivezProbe type and setLivezProbe() function to override the liveness probe with a test stub. test/mcp-standalone.test.ts imports these seams, defines a fast-fail probe mock, and verifies deterministic probe behavior without real network calls.

Changes

Probe Injection and Test Stabilization

Layer / File(s)	Summary
Probe injection seam src/mcp/rest-proxy.ts	Introduces exported LivezProbe type and default fetch-based implementation. New exported setLivezProbe(fn?) function allows tests to override the probe; resetHandleForTests() now resets the probe override to the default. Internal probe(url) calls the overridable livezProbe instead of fetch directly, with adjusted logging for optional status and statusText fields.
Test mocking and deterministic assertions test/mcp-standalone.test.ts	Imports resetHandleForTests and setLivezProbe from REST proxy. Defines instantLocalFallbackProbe mock returning ok:false immediately and fetchTrap that throws on unexpected calls. Test suite beforeEach resets the handle, installs the probe stub, and overrides globalThis.fetch; afterEach restores fetch and resets handle. New assertion test verifies handleToolCall("memory_save", ...) invokes the probe exactly once without triggering the fetch trap.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• rohitg00/agentmemory#161: Introduced the original /agentmemory/livez probe logic; this PR makes that probe injectable via the new seam.
• rohitg00/agentmemory#280: Also modifies REST proxy livez probe behavior; main PR adds injection capability while that PR adjusts timeout and environment-driven fallback control.

Poem

🐰 The probe once fetched from networks far,
Causing flakes beneath CI's star,
Now stubbed and swift, it fails on cue,
Tests run stable, green and true! ✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 25.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly and specifically describes the main change: fixing a test flake by injecting the livez probe, with direct reference to issue #449.
Linked Issues check	✅ Passed	The PR implements all core requirements from #449: module-level DI seam for livez probe via setLivezProbe(), deterministic test stubs avoiding real network calls, and regression test asserting proper probe invocation.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to solving the test flake: export of LivezProbe type and setLivezProbe function, test utilities integration, and probe stub setup in mcp-standalone.test.ts.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/449-mcp-standalone-probe-di

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 ESLint

If the error stems from missing dependencies, add them to the package.json file. For unrecoverable errors (e.g., due to private dependencies), disable the tool in the CodeRabbit configuration.

ESLint skipped: no ESLint configuration detected in root package.json. To enable, add eslint to devDependencies.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

src/mcp/rest-proxy.ts (1)

43-50: ⚡ Quick win

Remove WHAT-style explanatory comments in source per repo rule.

These blocks are explanatory comments about behavior; prefer self-descriptive naming and keep source comments minimal.

Proposed cleanup

-/**
- * Probes the agentmemory server's livez endpoint. Returns a Response-shaped
- * object whose `ok` flag drives the proxy/local-fallback decision.
- *
- * Tests can swap this via {`@link` setLivezProbe} to avoid the real 2s
- * AbortController race that destabilises mcp-standalone test runs (`#449`).
- * Production callers should leave it on the default.
- */
 export type LivezProbe = (
   url: string,
   timeoutMs: number,
   headers: Record<string, string>,
 ) => Promise<{ ok: boolean; status?: number; statusText?: string }>;
@@
-/**
- * Override the livez probe. Intended for tests — production code should rely
- * on the default fetch-based probe. Calling without an argument restores the
- * default. Pair with {`@link` resetHandleForTests} so the cached handle is
- * dropped before the next call.
- */
 export function setLivezProbe(fn?: LivezProbe): void {
   livezProbe = fn ?? defaultLivezProbe;
 }

As per coding guidelines, "Do not use code comments explaining WHAT — use clear naming instead".

Also applies to: 68-73

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/mcp/rest-proxy.ts` around lines 43 - 50, Replace the verbose WHAT-style
block comment with a concise, self-descriptive doc and/or a clearer name: remove
the multi-line explanatory comment above the livez probe implementation and
instead use a single-line JSDoc like "Default livez probe used by setLivezProbe"
or rename the probe to a descriptive identifier (e.g., defaultLivezProbe or
probeAgentMemoryLivez) so callers and tests (setLivezProbe) are
self-explanatory; apply the same cleanup to the other explanatory block at
68-73. Ensure the symbols setLivezProbe and the probe function/variable are
updated to match any renaming and keep comments minimal.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@src/mcp/rest-proxy.ts`:
- Around line 43-50: Replace the verbose WHAT-style block comment with a
concise, self-descriptive doc and/or a clearer name: remove the multi-line
explanatory comment above the livez probe implementation and instead use a
single-line JSDoc like "Default livez probe used by setLivezProbe" or rename the
probe to a descriptive identifier (e.g., defaultLivezProbe or
probeAgentMemoryLivez) so callers and tests (setLivezProbe) are
self-explanatory; apply the same cleanup to the other explanatory block at
68-73. Ensure the symbols setLivezProbe and the probe function/variable are
updated to match any renaming and keep comments minimal.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f70312b7-73d6-4a16-b61c-32ed8af26cf5

📥 Commits

Reviewing files that changed from the base of the PR and between c93c715 and 597ab47.

📒 Files selected for processing (2)

• src/mcp/rest-proxy.ts
• test/mcp-standalone.test.ts