fix(embedding): guard provider responses against dimension mismatches

[AmmarSaleh50]

Closes #247.

What

Adds a dimension-check wrapper at the EmbeddingProvider boundary in src/providers/embedding/index.ts. createEmbeddingProvider() and createImageEmbeddingProvider() now wrap every provider so that any returned Float32Array whose length differs from provider.dimensions throws a descriptive error.

Why

Embedding providers in src/providers/embedding/ (gemini, openai, voyage, cohere, openrouter, local, clip) trust that the API returns vectors matching their declared dimensions. None of them validate result.length === this.dimensions. When the assumption breaks, the failure is silent:

• src/state/vector-index.ts:10 returns 0 from cosineSimilarity on length mismatch instead of throwing.
• A wrong-size vector gets stored, never matches anything in search, and the affected memory becomes effectively invisible — no error, no log line.

This came up during review of #246 (gemini deprecation), where CodeRabbit flagged the gap on a single provider. Per-provider guards would be 7× duplication and easy to miss for new providers, so the fix lives at the factory boundary instead.

Changes

src/providers/embedding/index.ts

• New withDimensionGuard(provider) that returns a wrapper checking embed, embedBatch, and (when present) embedImage results.
• createEmbeddingProvider() and createImageEmbeddingProvider() apply the wrapper to every constructed provider.
• Errors name the provider, call site (embed / embedBatch[i] / embedImage), expected dimension, and actual dimension.

test/embedding-provider.test.ts

• Happy path passes through correct-dimension vectors.
• Mismatch on embed throws.
• Mismatch on any item in embedBatch throws with the index.
• embedImage is guarded when present and absent when not.

How to verify

npm install
npm test -- test/embedding-provider.test.ts

The four new tests under describe(\"withDimensionGuard\") should pass alongside the existing suite.

Notes

• Single signed-off commit per CONTRIBUTING.
• No CHANGELOG touch (release PRs only).
• No new dependencies.
• No existing test breakage expected — the existing tests construct providers directly via new GeminiEmbeddingProvider(...), which is unwrapped, while only the factory entry points apply the guard.

Summary by CodeRabbit

• New Features

  • Embedding providers now validate returned vector dimensions and surface clear errors (expected vs actual, including failing batch index).
  • Loading persisted vector indexes now verifies dimensionality and will either drop stale vectors with a warning (configurable) or refuse startup with guidance to re-embed or switch providers.
  • Indexes can report all observed stored vector dimensions.

• Tests

  • Added tests covering provider dimension guards and index dimension validation.

[vercel]

@AmmarSaleh50 is attempting to deploy a commit to the rohitg00's projects Team on Vercel.

A member of the Team first needs to authorize it.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds exported withDimensionGuard(provider) that enforces provider.dimensions on embed/embedBatch/(optional)embedImage outputs and wraps embedding factories; adds VectorIndex.validateDimensions(expected) and a restore-time check that compares persisted vectors to the active provider, optionally discarding or failing on mismatches; includes tests for the guard and index validation.

Changes

Embedding Dimension Validation & Index Safety

Layer / File(s)	Summary
VectorIndex: validateDimensions src/state/vector-index.ts	Add VectorIndex.validateDimensions(expected) which scans stored vectors, returns all mismatches (obsId and observed dim) and the set of distinct seen embedding dimensions.
Persisted Index Restore src/index.ts	On persisted vector index load, compute activeDim from embeddingProvider?.dimensions; when activeDim>0 call loaded.vector.validateDimensions(activeDim) and either discard persisted vectors when AGENTMEMORY_DROP_STALE_INDEX==="true" (warn) or throw an error refusing startup with mismatch/sample details; otherwise restore the index.
Guard Implementation src/providers/embedding/index.ts	Export withDimensionGuard(provider) that captures provider.dimensions and returns a prototype-preserving wrapper (via Object.create(provider)) overriding embed, embedBatch, and optional embedImage to validate returned Float32Array lengths and throw on mismatches.
Provider Integration src/providers/embedding/index.ts	createImageEmbeddingProvider() and createEmbeddingProvider() now wrap ClipEmbeddingProvider and detected providers (Gemini, OpenAI, Voyage, Cohere, OpenRouter, Local) with withDimensionGuard(...) before caching/returning.
Test Imports test/embedding-provider.test.ts	Update imports to include the shared EmbeddingProvider type from ../src/types.js.
Guard Validation Tests test/embedding-provider.test.ts	New describe("withDimensionGuard") tests check correct-dimension vectors pass; mismatched embed() and embedBatch() throw errors (including failing batch index); embedImage is conditionally guarded; wrapper preserves prototype instanceof.
VectorIndex Tests test/vector-index-dimensions.test.ts	Add Vitest tests for VectorIndex.validateDimensions() covering empty index, all-matching vectors, mixed-dimension mismatches (report all), and all-nonmatching scenarios.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

Possibly related issues

• Vector index restore path bypasses dimension guard from #248 #256: Adds persisted index dimension checks and VectorIndex.validateDimensions to compare active provider dimensions on restore — directly related to this PR's index-restore safety changes.

Possibly related PRs

• rohitg00/agentmemory#189: Changes how provider dimensions are computed for OpenAIEmbeddingProvider; related because the guard relies on provider.dimensions.
• rohitg00/agentmemory#179: Added Clip provider and image provider factory; related because this PR wraps the Clip image provider with the dimension guard.

"I count each floated line with care,
A rabbit's paw checks every pair,
If lengths align, the path stays true—
Else re-embed, rebuild, or view. 🐇"

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 20.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 'fix(embedding): guard provider responses against dimension mismatches' directly and clearly describes the main change: adding dimension validation at the embedding provider boundary.
Linked Issues check	✅ Passed	The PR fully addresses issue #247 by implementing the proposed dimension-guard wrapper at the provider factory boundary, validating embeddings at write-time, and extending validation to the persistence/load path.
Out of Scope Changes check	✅ Passed	All changes are directly scoped to dimension validation: provider wrapping, validation methods, and load-time checks. No unrelated refactorings or features were introduced.

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

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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]

Actionable comments posted: 1

🧹 Nitpick comments (1)

src/providers/embedding/index.ts (1)

52-55: ⚡ Quick win

Remove the WHAT-style comment block and let naming carry intent

This block explains behavior in prose; repo guidelines for src/**/*.ts ask to avoid WHAT-comments.

As per coding guidelines: “Avoid code comments explaining WHAT — use clear naming instead”.

🤖 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/providers/embedding/index.ts` around lines 52 - 55, Delete the WHAT-style
prose block that describes the silent failure and instead make the code
self-describing: remove the comment and add/rename a boundary function in
vector-index.ts to validate dimensions (e.g., ensureDimensionsMatchOrThrow) or
rename cosineSimilarity to reflect its behavior (e.g.,
cosineSimilarityOrThrowOnDimensionMismatch), and call that validator/wrapper
from the embedding entrypoint so mismatched-length vectors throw instead of
silently returning 0; this keeps intent in names rather than a prose comment.

🤖 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.

Inline comments:
In `@src/providers/embedding/index.ts`:
- Around line 66-75: The wrapper currently constructs a plain object assigned to
wrapped which loses the provider's prototype and breaks instanceof checks (e.g.,
GeminiEmbeddingProvider, OpenAIEmbeddingProvider); instead create the wrapper
with the original prototype (use Object.create(provider)) so class identity is
preserved, then override/assign the name, dimensions, embed and embedBatch
properties while keeping embed/embedBatch calls wrapped with the existing
check(...) logic to validate outputs.

---

Nitpick comments:
In `@src/providers/embedding/index.ts`:
- Around line 52-55: Delete the WHAT-style prose block that describes the silent
failure and instead make the code self-describing: remove the comment and
add/rename a boundary function in vector-index.ts to validate dimensions (e.g.,
ensureDimensionsMatchOrThrow) or rename cosineSimilarity to reflect its behavior
(e.g., cosineSimilarityOrThrowOnDimensionMismatch), and call that
validator/wrapper from the embedding entrypoint so mismatched-length vectors
throw instead of silently returning 0; this keeps intent in names rather than a
prose comment.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 733047e6-120c-4fb4-8253-e272662741b7

📥 Commits

Reviewing files that changed from the base of the PR and between 989e9c0 and 72506ac.

📒 Files selected for processing (2)

• src/providers/embedding/index.ts
• test/embedding-provider.test.ts

[rohitg00]

This is a clean, well-thought fix. Reviewed end to end:

Approach is right.
Putting the guard at the factory boundary in createEmbeddingProvider() / createImageEmbeddingProvider() instead of duplicating across every concrete provider. New providers pick up the guard for free without anyone having to remember to add it.

Prototype-chain preservation is the detail that matters.
Object.create(provider) so instanceof GeminiEmbeddingProvider keeps working through the wrapper. That's the bit most "wrap-and-forward" implementations get wrong, and it would have bit users who introspect the provider downstream.

Per-vector check on embedBatch with the indexed error message (embedBatch[3]) is exactly what you want when debugging — points straight at the bad vector in a 100-element batch instead of just "something is wrong."

Test coverage is comprehensive. Good and bad paths for embed, embedBatch, and embedImage. The fake-provider helper is the right shape for this.

Together with #246 (Gemini migration that explicitly passes outputDimensionality), these two PRs close the silent-corruption window that #247 documents. If Google ever ignores the dim param on a later API change, #248 throws loudly instead of polluting the index with zero-similarity vectors.

No blockers. Approving — also in the land-soon bucket, holding for @rohitg00.

Two things to consider after merge (not blocking):

1. CHANGELOG entry under Fixed should call out both Embedding providers silently corrupt the vector index when an API returns wrong-dimension vectors #247 and the linkage to fix(embedding): migrate Gemini provider to gemini-embedding-001 #246 so users understand the combined fix.
2. Worth a "lessons learned" note in the embedding-provider doc / contributor guide: every new provider's dimensions field is a contract the wrapper now enforces. Saves the next contributor from learning this the hard way.

[rohitg00]

Re-reviewed independently against the alternatives we'd consider, not just rubber-stamping. Conclusion: this is the right approach. Specifically:

Why factory-boundary > the alternatives

Approach	Why we'd reject
Per-provider validation in each embed/embedBatch	Duplicated across 7 providers, easy to forget when adding the 8th. Footgun lives forever.
Validate at vectorIndex.add()	Pushes responsibility to every caller. Embedding flows through observe, replay-import, restore — each would need its own check.
Fix cosineSimilarity to throw instead of returning 0	Catches at search time, after bad data is in the index. Re-embedding to recover takes hours on large stores per #247.
Factory wrapper (this PR)	Single point of truth, runs at the moment data crosses from "untrusted external API" to "internal index".

The Object.create(provider) to preserve the prototype chain is the detail most "wrap and forward" implementations miss. Verified by your test that instanceof FakeProvider still resolves through the wrapper — that matters for any downstream code that introspects which concrete provider is in use.

The per-vector index in embedBatch[i] error message is the kind of thing that pays for itself the first time someone debugs a 100-element batch failure.

Together with #246, your two PRs close the silent-corruption window front-to-back: #246 explicitly passes outputDimensionality so providers are asked for the right size; #248 throws if they ever ignore that and return the wrong size anyway.

One real follow-up gap (not blocking this PR):

The IndexPersistence.load() path at src/state/index-persistence.ts:51-66 deserializes vectors directly from KV with no dimension check. If a user persists a 384-dim MiniLM index and then upgrades to 1536-dim OpenAI without re-indexing, the restore brings 384-dim vectors back into a freshly-instantiated VectorIndex, and live observations from the new provider write 1536-dim ones alongside. cosineSimilarity still returns 0 on the cross-dim comparisons, silent recall degradation again — same failure mode this PR fixes for the live-API path, just via a different on-ramp.

Worth a follow-up issue: validate dimensions against the active provider at IndexPersistence.load() time, and either re-index or refuse to start if they mismatch. Happy to file the follow-up after this lands so it's traceable from the same audit thread.

Verdict: APPROVE. This was the highest-quality PR in the current backlog — issue + fix + tests in one well-scoped contribution. Strong work, @AmmarSaleh50.

Holding for @rohitg00 to push the merge button.

[rohitg00]

Pushed 57311c6 to your branch (since maintainerCanModify: true) to close the restore-path gap I flagged in the review. This stays inside the same PR so the live-write fix and its symmetric persistence fix land together — easier to review, easier to revert as a unit, and the CHANGELOG entry covers both on-ramps.

What the commit adds

1. VectorIndex.firstDimensions() — returns the dimension of any stored vector (or 0 if empty). All vectors in a single index share a dimension because the live-write path is now gated by your wrapper.

2. Dimension validation at startup in src/index.ts after IndexPersistence.load() returns. When the persisted index dimension differs from the active provider's dimensions:

• Default: refuse to start with a clear, actionable error:

  [agentmemory] Refusing to start: persisted vector index has dimension 384,
  but the active provider (openai) declares 1536. Loading would silently corrupt
  search (cross-dimension cosine returns 0). Choose one:
    - Re-embed the existing index against the new provider, then start.
    - Set AGENTMEMORY_DROP_STALE_INDEX=true to discard the persisted vectors
      and rebuild from live observations.
    - Switch the embedding provider back to the one that wrote the index.
  

• Opt-out: AGENTMEMORY_DROP_STALE_INDEX=true logs a warning, discards the persisted vectors, lets the live path rebuild. Friendlier for users who deliberately swap providers.

3. Test in test/vector-index-dimensions.test.ts covering the three cases (empty, small, large 1536-dim).

Why it stays in scope

• Same root cause (Embedding providers silently corrupt the vector index when an API returns wrong-dimension vectors #247): cosineSimilarity returns 0 on length mismatch instead of throwing.
• Same blast radius: silent recall corruption, no log line, requires re-embedding to recover.
• Same defense pattern: refuse to let bad data into the index in the first place.
• Different on-ramp: yours catches fresh API responses; this catches persisted vectors loaded after a config change.

#256 is closed by this commit too — should auto-close on merge.

Tested locally:

• npx tsc --noEmit shows no new errors from these files.
• npx vitest run test/vector-index-dimensions.test.ts — 3/3 pass.
• Existing test/embedding-provider.test.ts from your earlier commits still passes (your withDimensionGuard tests untouched).

If you'd rather have this as a separate PR for clean attribution / lighter review surface, say the word and I'll split it. But I think the bundle is more useful here.

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (2)

src/index.ts (1)

339-345: ⚡ Quick win

Trim explanatory inline comments and rely on the guard logic + naming.

This block can stay clear without the added “what” comments.

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

🤖 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/index.ts` around lines 339 - 345, Remove the verbose explanatory inline
comments in the vector-dimension guard block and replace them with a concise
comment that names the guard behavior; keep a short note like "Refuse to load
vectors with mismatched dimension." and rely on the existing guard logic (the
dimension check that refuses loading and the same defense used by the live-write
path referenced in the code) and symbols such as cosineSimilarity to make intent
clear; do not change the guard logic itself—only trim the explanatory text to a
single-line summary.

src/state/vector-index.ts (1)

69-72: ⚡ Quick win

Remove inline “what” comments and keep this self-descriptive via naming.

These comments describe behavior directly and can be removed to align with the project convention.

♻️ Proposed cleanup

-  // Dimension of any stored vector, or 0 if the index is empty. All vectors
-  // in a single index are expected to share the same dimension; the first
-  // entry is representative because the live-write path is gated by the
-  // dimension guard in src/providers/embedding/index.ts.
   firstDimensions(): number {

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

🤖 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/state/vector-index.ts` around lines 69 - 72, Remove the inline "what"
comment block describing vector dimensions and make the code self-descriptive:
delete the three-line comment and, if needed, rename the exposed identifier
(e.g., dimension or vectorDimension/embeddingDimension) or add a clearer
identifier so the comment's information is conveyed by the name of the
variable/property in the vector index module (e.g., vectorDimension or
embeddingDimension) and any brief docstring can be limited to why rather than
what.

🤖 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.

Inline comments:
In `@src/index.ts`:
- Around line 346-348: The code currently uses loaded.vector.firstDimensions()
(persistedDim) to compare against embeddingProvider?.dimensions (activeDim),
which misses mixed-dimension entries; update the validation to iterate over all
persisted vectors in loaded.vector and verify each vector's dimension equals
activeDim (or fail/skip any that don't) instead of relying on firstDimensions();
apply the same fix to the other similar check around the second occurrence
(lines referenced in the review). Ensure you log or surface which vector(s) were
invalid and prevent restoring indexes containing mismatched-dimension vectors.

---

Nitpick comments:
In `@src/index.ts`:
- Around line 339-345: Remove the verbose explanatory inline comments in the
vector-dimension guard block and replace them with a concise comment that names
the guard behavior; keep a short note like "Refuse to load vectors with
mismatched dimension." and rely on the existing guard logic (the dimension check
that refuses loading and the same defense used by the live-write path referenced
in the code) and symbols such as cosineSimilarity to make intent clear; do not
change the guard logic itself—only trim the explanatory text to a single-line
summary.

In `@src/state/vector-index.ts`:
- Around line 69-72: Remove the inline "what" comment block describing vector
dimensions and make the code self-descriptive: delete the three-line comment
and, if needed, rename the exposed identifier (e.g., dimension or
vectorDimension/embeddingDimension) or add a clearer identifier so the comment's
information is conveyed by the name of the variable/property in the vector index
module (e.g., vectorDimension or embeddingDimension) and any brief docstring can
be limited to why rather than what.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 3ff42410-6f54-4f6f-af90-3e3af81c357d

📥 Commits

Reviewing files that changed from the base of the PR and between 35bb2d5 and 57311c6.

📒 Files selected for processing (3)

• src/index.ts
• src/state/vector-index.ts
• test/vector-index-dimensions.test.ts

[rohitg00]

Pushed b37f9ed to address the review feedback on 57311c6.

Verification per finding

Finding	Status	Action
firstDimensions() misses mixed-dimension entries inside a single persisted index	Valid	Fixed
"Apply the same fix to the other similar check around the second occurrence"	Not valid for this codebase	Skipped — see below

The only restoreFrom site besides vectorIndex.restoreFrom is bm25Index.restoreFrom at src/index.ts:333. BM25 entries don't carry dimensions, so there's no analogous check to update. Confirmed by grep: firstDimensions had exactly one caller, restoreFrom has two callers across the file (one BM25, one vector).

What changed

• VectorIndex.firstDimensions() → VectorIndex.validateDimensions(expected). New method walks every stored vector and returns:

  { mismatches: Array<{ obsId: string; dim: number }>; seenDimensions: Set<number> }

  No more "trust the first entry" — a legacy on-disk index that mixes dimensions (which used to be possible before fix(embedding): guard provider responses against dimension mismatches #248 added the live-API guard) is now caught.

• src/index.ts uses the new method. The startup error now reports:

  • How many of the total vectors are wrong (123 of 5000 vectors)
  • The distinct dimensions actually present on disk (dimensions seen on disk: 384, 1536)
  • The first 5 mismatched obsIds (First mismatched obsIds: o42 (dim=384), o91 (dim=1536), ...)

  So a debugging operator can spot-check what's in their index without needing to dump it.

• Test rewritten to cover the case firstDimensions() would have missed:

  it("reports every wrong-dimension vector, not just the first", ...);

  Plus empty / all-match / entire-index-mismatch for completeness.

Validation

• npx tsc --noEmit — no errors from the changed files.
• npx vitest run test/vector-index-dimensions.test.ts test/embedding-provider.test.ts — 21/21 pass (4 new + 17 existing from the live-write guard tests).

PR is now ready for merge from my side.

[coderabbitai]

🧹 Nitpick comments (1)

src/state/vector-index.ts (1)

69-76: ⚡ Quick win

Remove the explanatory WHAT block comment and rely on naming/types.

This comment narrates behavior rather than encoding intent in code, which conflicts with the repo rule for src/**/*.ts.

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

🤖 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/state/vector-index.ts` around lines 69 - 76, Remove the multi-line
explanatory WHAT comment above the vector-walking logic and rely on clear naming
and types instead: delete that block and ensure the surrounding symbols (the
function that walks stored vectors, and its return fields/variables named
mismatches, seenDimensions, and expected) have descriptive names and proper
TypeScript typings so the behavior is obvious without the long narrative
comment.

🤖 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/state/vector-index.ts`:
- Around line 69-76: Remove the multi-line explanatory WHAT comment above the
vector-walking logic and rely on clear naming and types instead: delete that
block and ensure the surrounding symbols (the function that walks stored
vectors, and its return fields/variables named mismatches, seenDimensions, and
expected) have descriptive names and proper TypeScript typings so the behavior
is obvious without the long narrative comment.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 25263d18-d616-414f-835a-3bb694367c83

📥 Commits

Reviewing files that changed from the base of the PR and between 57311c6 and b37f9ed.

📒 Files selected for processing (3)

• src/index.ts
• src/state/vector-index.ts
• test/vector-index-dimensions.test.ts

🚧 Files skipped from review as they are similar to previous changes (1)

• src/index.ts