fix(firestore-vector-search): restore terminal-state skip on embedOnWrite

kits/firestore-vector-search/README.md

@@ -251,23 +251,14 @@ longer written, only `COMPLETED` and `ERROR`. Anything reading
 `status.<instance id>.state`, or a security rule or index keyed to it, needs
 updating. The field name is still `STATUS_FIELD_NAME`, defaulting to `status`.
 
+`embedOnWrite` still reads the same four states the extension treated as final —
+`PROCESSING`, `COMPLETED`, `ERROR` and `BACKFILLED` — so documents an installed
+instance already embedded are still skipped once you flatten their status field.
+
 Query documents no longer get a status field at all. They previously carried
 `status.textQuery`, so if you were waiting on that to know a query had finished,
 wait for `result` instead.
 
-### Editing a document's input re-embeds it
-
-The extension embedded each document once. Its skip rule was "this document's
-status is already in a final state", so once a document reached `COMPLETED` (or
-`ERROR`), changing its input field never produced a new embedding and a failure
-was never retried.
-
-The kit compares the input instead: it re-embeds when the input field changes, and
-skips only when the input is unchanged and an embedding is already present. This
-is usually what you wanted, but it means editing inputs in bulk now costs
-embedding calls, and a document that previously sat stale will be brought up to
-date on its next write.
-
 ### The lifecycle hooks and the function region
 
 Install and reconfigure hooks are replaced by an `initVectorSearch` task that the

kits/firestore-vector-search/src/handlers.ts

@@ -47,6 +47,27 @@ export type VectorWriteEvent = FirestoreEvent<
   Record<string, string>
 >;
 
+/**
+ * States the extension's `FirestoreOnWriteProcessor` treated as final. A
+ * document that has reached one of these is never processed again, so each
+ * document is embedded once and a failure is never retried.
+ */
+const TERMINAL_STATES = new Set([
+  "PROCESSING",
+  "COMPLETED",
+  "ERROR",
+  "BACKFILLED",
+]);
+
+function isInTerminalState(
+  data: FirebaseFirestore.DocumentData,
+  statusFieldName: string
+): boolean {
+  const status = data[statusFieldName] as { state?: unknown } | undefined;
+  const state = status?.state;
+  return typeof state === "string" && TERMINAL_STATES.has(state);
+}
+
 function queuePath(
   config: ResolvedVectorSearchConfig,
   queueName: string
@@ -78,12 +99,9 @@ export async function handleEmbedOnWrite(
   logs.start("embedOnWrite");
 
   const data = event.data.after.data() ?? {};
+  if (isInTerminalState(data, ctx.config.statusFieldName)) return;
   const input = data[ctx.config.inputFieldName];
   if (typeof input !== "string") return;
-  const beforeInput = event.data.before.exists
-    ? event.data.before.get(ctx.config.inputFieldName)
-    : undefined;
-  if (beforeInput === input && data[ctx.config.outputFieldName]) return;
 
   try {
     const embedding = await embedClient(ctx).getSingleEmbedding(input);

kits/firestore-vector-search/tests/handlers.test.ts

@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+import { FieldValue } from "firebase-admin/firestore";
 import type { CallableRequest } from "firebase-functions/v2/https";
 import { HttpsError } from "firebase-functions/v2/https";
 import { beforeEach, describe, expect, test, vi } from "vitest";
@@ -36,7 +37,12 @@ vi.mock("../src/embeddings", () => ({
 // handler never needs it.
 vi.mock("../src/queries/setup", () => ({ createIndex: vi.fn() }));
 
-import { type HandlerContext, handleQueryCall } from "../src/handlers";
+import {
+  type HandlerContext,
+  type VectorWriteEvent,
+  handleEmbedOnWrite,
+  handleQueryCall,
+} from "../src/handlers";
 import { resolveVectorSearchConfig } from "../src/export-config";
 
 const config = resolveVectorSearchConfig({
@@ -205,3 +211,184 @@ describe("handleQueryCall", () => {
     expect((err as Error).message).toBe("Query failed");
   });
 });
+
+/** A minimal `DocumentSnapshot` stand-in backed by a plain object. */
+function snapshot(data: Record<string, unknown> | null) {
+  const set = vi.fn().mockResolvedValue(undefined);
+  return {
+    snap: {
+      exists: data !== null,
+      data: () => data ?? undefined,
+      get: (field: string) => data?.[field],
+      ref: { path: `${config.collectionPath}/doc-1`, set },
+    },
+    set,
+  };
+}
+
+function writeEvent(
+  before: Record<string, unknown> | null,
+  after: Record<string, unknown> | null
+) {
+  const beforeSnap = snapshot(before);
+  const afterSnap = snapshot(after);
+  const event = {
+    params: { docId: "doc-1" },
+    data: { before: beforeSnap.snap, after: afterSnap.snap },
+  } as unknown as VectorWriteEvent;
+  return { event, set: afterSnap.set };
+}
+
+function embedCtx() {
+  return { firestore: {}, config } as unknown as HandlerContext;
+}
+
+describe("handleEmbedOnWrite", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    getSingleEmbedding.mockResolvedValue(EMBEDDING);
+  });
+
+  test("embeds a new document and marks it COMPLETED", async () => {
+    const { event, set } = writeEvent(null, { input: "hello" });
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).toHaveBeenCalledWith("hello");
+    expect(set).toHaveBeenCalledWith(
+      {
+        [config.outputFieldName]: FieldValue.vector(EMBEDDING),
+        [config.statusFieldName]: { state: "COMPLETED" },
+      },
+      { merge: true }
+    );
+  });
+
+  test("embeds a document that already has an embedding but no status", async () => {
+    const { event } = writeEvent(null, {
+      input: "hello",
+      [config.outputFieldName]: FieldValue.vector(EMBEDDING),
+    });
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).toHaveBeenCalledWith("hello");
+  });
+
+  test("marks the document ERROR and rethrows when embedding fails", async () => {
+    const { event, set } = writeEvent(null, { input: "hello" });
+    getSingleEmbedding.mockRejectedValue(new Error("Embedding failed"));
+
+    await expect(handleEmbedOnWrite(event, embedCtx())).rejects.toThrow(
+      "Embedding failed"
+    );
+    expect(set).toHaveBeenCalledWith(
+      {
+        [config.statusFieldName]: {
+          state: "ERROR",
+          message: "Embedding failed",
+        },
+      },
+      { merge: true }
+    );
+  });
+
+  test("skips a deleted document", async () => {
+    const { event, set } = writeEvent({ input: "hello" }, null);
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).not.toHaveBeenCalled();
+    expect(set).not.toHaveBeenCalled();
+  });
+
+  test("skips a document whose input is not a string", async () => {
+    const { event, set } = writeEvent(null, { input: 42 });
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).not.toHaveBeenCalled();
+    expect(set).not.toHaveBeenCalled();
+  });
+
+  // Parity with the extension: `FirestoreOnWriteProcessor` skipped any document
+  // already in a final state, so an edited input never produced a new embedding
+  // and a failure was never retried.
+  for (const state of ["PROCESSING", "COMPLETED", "ERROR", "BACKFILLED"]) {
+    test(`does not re-embed a document in the ${state} state`, async () => {
+      const { event, set } = writeEvent(
+        { input: "hello", [config.statusFieldName]: { state } },
+        { input: "goodbye", [config.statusFieldName]: { state } }
+      );
+
+      await handleEmbedOnWrite(event, embedCtx());
+
+      expect(getSingleEmbedding).not.toHaveBeenCalled();
+      expect(set).not.toHaveBeenCalled();
+    });
+  }
+
+  test("embeds a document in an unrecognised state", async () => {
+    const { event } = writeEvent(null, {
+      input: "hello",
+      [config.statusFieldName]: { state: "SOMETHING_ELSE" },
+    });
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).toHaveBeenCalledWith("hello");
+  });
+
+  // The extension's skip rule was the status state alone, with no comparison
+  // against the previous input, so an unchanged document with an embedding but
+  // no status was still processed.
+  test("embeds an unchanged document that has an embedding but no status", async () => {
+    const doc = {
+      input: "hello",
+      [config.outputFieldName]: FieldValue.vector(EMBEDDING),
+    };
+    const { event } = writeEvent({ ...doc }, { ...doc });
+
+    await handleEmbedOnWrite(event, embedCtx());
+
+    expect(getSingleEmbedding).toHaveBeenCalledWith("hello");
+  });
+
+  describe("with a custom status field name", () => {
+    const customConfig = resolveVectorSearchConfig({
+      projectId: "test-project",
+      instanceId: "test-instance",
+      statusFieldName: "embedStatus",
+    });
+
+    function customCtx() {
+      return {
+        firestore: {},
+        config: customConfig,
+      } as unknown as HandlerContext;
+    }
+
+    test("skips on the configured field", async () => {
+      const { event, set } = writeEvent(null, {
+        input: "hello",
+        embedStatus: { state: "COMPLETED" },
+      });
+
+      await handleEmbedOnWrite(event, customCtx());
+
+      expect(getSingleEmbedding).not.toHaveBeenCalled();
+      expect(set).not.toHaveBeenCalled();
+    });
+
+    test("ignores a terminal state on the default field", async () => {
+      const { event } = writeEvent(null, {
+        input: "hello",
+        status: { state: "COMPLETED" },
+      });
+
+      await handleEmbedOnWrite(event, customCtx());
+
+      expect(getSingleEmbedding).toHaveBeenCalledWith("hello");
+    });
+  });
+});