From e0eccbdd7e2094f97f32e67fdf1da71f80a97964 Mon Sep 17 00:00:00 2001
From: kayo09 <68217041+kayo09@users.noreply.github.com>
Date: Sat, 4 Apr 2026 18:24:12 +0530
Subject: [PATCH] fix(lib): guard Unlimiformer attentionScore against
 non-finite embeddings

Return NaN when vectors contain NaN, Infinity, or non-numbers instead of throwing from cosineSimilarity. Skip non-finite query and item embeddings in rankItemsByAttentionTopK. Add tests.
---
 packages/lib/unlimiformer.test.ts | 41 +++++++++++++++++++++++++++++++
 packages/lib/unlimiformer.ts      | 21 +++++++++++++---
 2 files changed, 58 insertions(+), 4 deletions(-)

diff --git a/packages/lib/unlimiformer.test.ts b/packages/lib/unlimiformer.test.ts
index cd72921e..3c134d9f 100644
--- a/packages/lib/unlimiformer.test.ts
+++ b/packages/lib/unlimiformer.test.ts
@@ -1,5 +1,6 @@
 import { describe, expect, test } from "bun:test"
 import {
+	attentionScore,
 	attentionScores,
 	rankItemsByAttentionTopK,
 	topKAttentionKeys,
@@ -60,6 +61,28 @@ describe("topKAttentionKeys", () => {
 			),
 		).toEqual([])
 	})
+
+	test("skips keys with NaN components without throwing", () => {
+		const q = unit(1, 0, 0)
+		const keys = [unit(1, 0, 0), [1, Number.NaN, 0], unit(0, 1, 0)]
+		const top = topKAttentionKeys(q, keys, 5)
+		expect(top.map((t) => t.index)).toEqual([0, 2])
+	})
+})
+
+describe("attentionScore", () => {
+	test("returns NaN instead of throwing when a vector contains NaN", () => {
+		const k = unit(1, 0, 0)
+		expect(attentionScore([1, Number.NaN, 0], k)).toBeNaN()
+		expect(attentionScore(k, [1, Number.NaN, 0])).toBeNaN()
+	})
+
+	test("returns NaN for non-number or non-finite components", () => {
+		const k = unit(1, 0, 0)
+		const stringSlot = [1, "x", 0] as unknown as number[]
+		expect(attentionScore(stringSlot, k)).toBeNaN()
+		expect(attentionScore([Number.POSITIVE_INFINITY, 0, 0], k)).toBeNaN()
+	})
 })
 
 describe("attentionScores", () => {
@@ -113,4 +136,22 @@ describe("rankItemsByAttentionTopK", () => {
 		expect(ranked).toHaveLength(1)
 		expect(ranked[0]?.item.id).toBe("ok")
 	})
+
+	test("returns empty when query embedding is non-finite", () => {
+		const items = [{ id: "a", e: unit(1, 0, 0) }]
+		expect(
+			rankItemsByAttentionTopK([Number.NaN, 0, 0], items, (x) => x.e, 2),
+		).toEqual([])
+	})
+
+	test("skips items with non-finite embeddings", () => {
+		const items = [
+			{ id: "bad", e: [1, Number.NaN, 0] },
+			{ id: "ok", e: unit(0, 1, 0) },
+		]
+		const q = unit(0, 1, 0)
+		const ranked = rankItemsByAttentionTopK(q, items, (x) => x.e, 2)
+		expect(ranked).toHaveLength(1)
+		expect(ranked[0]?.item.id).toBe("ok")
+	})
 })
diff --git a/packages/lib/unlimiformer.ts b/packages/lib/unlimiformer.ts
index 882fae9c..5403c67f 100644
--- a/packages/lib/unlimiformer.ts
+++ b/packages/lib/unlimiformer.ts
@@ -9,6 +9,10 @@
 
 import { cosineSimilarity } from "./similarity"
 
+/** True when every entry is a finite number (empty arrays allowed). */
+const isFiniteEmbeddingVector = (v: number[]): boolean =>
+	v.every((x) => typeof x === "number" && Number.isFinite(x))
+
 export type AttentionTopK = {
 	index: number
 	score: number
@@ -19,19 +23,23 @@ export type AttentionTopK = {
  * For normalized embeddings this matches cosine similarity.
  *
  * Returns `NaN` when `query` and `key` have different lengths (e.g. mixed embedding
- * models) so callers can avoid throwing from `cosineSimilarity`.
+ * models), or when either vector contains non-finite values (`NaN`, `±Infinity`), so
+ * callers avoid throwing from `cosineSimilarity`.
  */
 export const attentionScore = (query: number[], key: number[]): number => {
 	if (query.length !== key.length) {
 		return Number.NaN
 	}
+	if (!isFiniteEmbeddingVector(query) || !isFiniteEmbeddingVector(key)) {
+		return Number.NaN
+	}
 	return cosineSimilarity(query, key)
 }
 
 /**
  * Attention scores for `query` against every row in `keys`, aligned by index.
- * Entries are `NaN` when a key length does not match the query (same embedding model
- * is required for a meaningful score).
+ * Entries are `NaN` when a key length does not match the query, or when either vector
+ * has non-finite components.
  */
 export const attentionScores = (query: number[], keys: number[][]): number[] =>
 	keys.map((key) => attentionScore(query, key))
@@ -97,6 +105,10 @@ export const rankItemsByAttentionTopK = <T>(
 		return []
 	}
 
+	if (!isFiniteEmbeddingVector(queryEmbedding)) {
+		return []
+	}
+
 	const packed: Array<{ item: T; originalIndex: number; embedding: number[] }> =
 		[]
 
@@ -107,7 +119,8 @@ export const rankItemsByAttentionTopK = <T>(
 		if (
 			embedding &&
 			embedding.length > 0 &&
-			embedding.length === queryEmbedding.length
+			embedding.length === queryEmbedding.length &&
+			isFiniteEmbeddingVector(embedding)
 		) {
 			packed.push({ item, originalIndex: i, embedding })
 		}