From 587e65500e60dbb08be90e10a6e98275ee910296 Mon Sep 17 00:00:00 2001 From: Corie Watson Date: Wed, 2 Sep 2026 15:01:01 +0100 Subject: [PATCH 1/2] fix(firestore-vector-search): guard queryOnWrite with the extension's status handleQueryOnWrite wrote its result back onto the document it watches with no status or idempotency guard, so the write re-fired the trigger and the handler queried again, billing one embedding call plus one vector query per pass. The extension ran the query as a FirestoreOnWriteProcess through a FirestoreOnWriteProcessor, which skipped any document whose status state was already PROCESSING, COMPLETED, ERROR or BACKFILLED and only ran when one of the fieldDependencyArray fields (query, limit) changed. That processor is firebase-functions v1-tied so the kit cannot consume it, but its behaviour ports directly: write status.textQuery PROCESSING before the query, COMPLETED (or ERROR) with the result after it, and skip on a set state or unchanged inputs. The PROCESSING write means neither status write can re-trigger a run. The processor was built without a statusField, so query documents used its literal `status` default rather than STATUS_FIELD_NAME; that is kept as-is. A failed query now marks the document ERROR and the invocation succeeds instead of throwing, matching the extension's errorFn. Fixes #3010 --- kits/firestore-vector-search/README.md | 12 +- kits/firestore-vector-search/src/handlers.ts | 81 ++++++- .../tests/handlers.test.ts | 219 +++++++++++++++++- 3 files changed, 298 insertions(+), 14 deletions(-) diff --git a/kits/firestore-vector-search/README.md b/kits/firestore-vector-search/README.md index 0552ca77b5..92d1d46b7c 100644 --- a/kits/firestore-vector-search/README.md +++ b/kits/firestore-vector-search/README.md @@ -251,9 +251,9 @@ longer written, only `COMPLETED` and `ERROR`. Anything reading `status..state`, or a security rule or index keyed to it, needs updating. The field name is still `STATUS_FIELD_NAME`, defaulting to `status`. -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. +This applies to the documents in your indexed collection only. Query documents +keep the extension's nested shape, `status.textQuery`, with the same states and +timestamps, so anything waiting on that path still works. ### Editing a document's input re-embeds it @@ -307,7 +307,11 @@ for; the Firebase CLI grants these for you. `status`, and embeddings are still written as native Firestore vectors. - Querying by writing a document to `_/index/queries` still works the same way, with `query`, an optional `limit` and optional `prefilters`, and - the matching document ids written back to the document under `result`. + the matching document ids written back to the document under `result`. A query + document is still run once: writes that leave `query` and `limit` untouched are + ignored, as is any write to a document whose `status.textQuery.state` is + already set, so editing a completed query document does not re-run it and a + failed one is not retried. Write a new document for a new query. - `queryCallable` still requires an authenticated caller, still validates its argument with the same schema, still rejects a `limit` that is not an integer above zero, and still returns `{ ids: [...] }`. diff --git a/kits/firestore-vector-search/src/handlers.ts b/kits/firestore-vector-search/src/handlers.ts index d62061ac9e..326a082310 100644 --- a/kits/firestore-vector-search/src/handlers.ts +++ b/kits/firestore-vector-search/src/handlers.ts @@ -117,25 +117,88 @@ export async function handleEmbedOnWrite( } } +/** + * The extension ran the query through a `FirestoreOnWriteProcessor` built with + * no `statusField`, so query documents used the processor's literal `status` + * default rather than `STATUS_FIELD_NAME`, under the process id `textQuery`. + */ +const QUERY_STATUS_PATH = "status.textQuery"; + +/** The fields the extension's `fieldDependencyArray` watched. */ +const QUERY_INPUT_FIELDS = ["query", "limit"] as const; + +/** + * States the extension's processor treated as final: a query document in any of + * them is never processed again. `PROCESSING` is one of them, so neither of the + * status writes below can re-trigger a run. + */ +const FINAL_QUERY_STATES = new Set([ + "PROCESSING", + "COMPLETED", + "ERROR", + "BACKFILLED", +]); + export async function handleQueryOnWrite( event: VectorWriteEvent, ctx: HandlerContext ): Promise { if (!event.data?.after.exists) return; - const data = event.data.after.data() ?? {}; + const after = event.data.after; + const data = after.data() ?? {}; const query = data.query; if (typeof query !== "string") return; - const result = await performTextQuery({ - query, - limit: data.limit ? parseLimit(data.limit) : ctx.config.defaultQueryLimit, - prefilters: (data.prefilters as Prefilter[] | undefined) ?? [], - embedClient: embedClient(ctx), - vectorStore: vectorStore(ctx), - config: ctx.config, + // The result write below re-fires this trigger. Skipping documents whose + // status is already final stops the loop, and matches the extension: a query + // document runs once and is never re-run, not even when its inputs change. + const state = after.get(`${QUERY_STATUS_PATH}.state`); + if (typeof state === "string" && FINAL_QUERY_STATES.has(state)) return; + + const before = event.data.before.exists + ? event.data.before.data() + : undefined; + const inputsChanged = QUERY_INPUT_FIELDS.some( + (field) => data[field] !== before?.[field] + ); + if (!inputsChanged) return; + + const startTime = FieldValue.serverTimestamp(); + await after.ref.update({ + [QUERY_STATUS_PATH]: { + state: "PROCESSING", + startTime, + createTime: + after.get(`${QUERY_STATUS_PATH}.createTime`) || after.createTime, + updateTime: startTime, + }, }); - await event.data.after.ref.set(result, { merge: true }); + try { + const result = await performTextQuery({ + query, + limit: data.limit ? parseLimit(data.limit) : ctx.config.defaultQueryLimit, + prefilters: (data.prefilters as Prefilter[] | undefined) ?? [], + embedClient: embedClient(ctx), + vectorStore: vectorStore(ctx), + config: ctx.config, + }); + const completeTime = FieldValue.serverTimestamp(); + await after.ref.update({ + ...result, + [`${QUERY_STATUS_PATH}.state`]: "COMPLETED", + [`${QUERY_STATUS_PATH}.updateTime`]: completeTime, + [`${QUERY_STATUS_PATH}.completeTime`]: completeTime, + }); + } catch (err) { + // The extension's `errorFn` logged and swallowed, so a failed query marks + // the document ERROR and the invocation succeeds rather than retrying. + logs.error("queryOnWrite", err); + await after.ref.update({ + [`${QUERY_STATUS_PATH}.state`]: "ERROR", + [`${QUERY_STATUS_PATH}.updateTime`]: FieldValue.serverTimestamp(), + }); + } } export async function handleQueryCall( diff --git a/kits/firestore-vector-search/tests/handlers.test.ts b/kits/firestore-vector-search/tests/handlers.test.ts index edcfc10245..ce60e34ec4 100644 --- a/kits/firestore-vector-search/tests/handlers.test.ts +++ b/kits/firestore-vector-search/tests/handlers.test.ts @@ -36,7 +36,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, + handleQueryCall, + handleQueryOnWrite, +} from "../src/handlers"; import { resolveVectorSearchConfig } from "../src/export-config"; const config = resolveVectorSearchConfig({ @@ -205,3 +210,215 @@ describe("handleQueryCall", () => { expect((err as Error).message).toBe("Query failed"); }); }); + +/** + * A DocumentSnapshot stand-in. `undefined` data means the snapshot does not + * exist, matching a create's `before` or a delete's `after`. + */ +function snapshot( + data: Record | undefined, + createTime: unknown = "doc-create-time" +) { + const update = vi.fn().mockResolvedValue(undefined); + return { + exists: data !== undefined, + createTime, + data: () => data, + get: (path: string) => + path + .split(".") + .reduce( + (acc, key) => + acc == null ? undefined : (acc as Record)[key], + data + ), + ref: { path: "queries/query-1", update }, + update, + }; +} + +function writeEvent( + before: ReturnType, + after: ReturnType +) { + return { + data: { before, after }, + params: { queryId: "query-1" }, + } as unknown as VectorWriteEvent; +} + +const COMPLETED_STATUS = { + status: { textQuery: { state: "COMPLETED" } }, +}; + +describe("handleQueryOnWrite", () => { + beforeEach(() => { + vi.clearAllMocks(); + getSingleEmbedding.mockResolvedValue(EMBEDDING); + }); + + test("runs the query on create and writes the result with a status", async () => { + const { ctx, chain } = makeCtx(); + const after = snapshot({ query: "test query" }); + + await handleQueryOnWrite(writeEvent(snapshot(undefined), after), ctx); + + expect(getSingleEmbedding).toHaveBeenCalledWith("test query"); + expect(chain.findNearest).toHaveBeenCalledWith( + config.outputFieldName, + EMBEDDING, + { + limit: config.defaultQueryLimit, + distanceMeasure: config.distanceMeasure, + } + ); + expect(after.update).toHaveBeenCalledTimes(2); + + const start = after.update.mock.calls[0][0]; + expect(start["status.textQuery"]).toMatchObject({ + state: "PROCESSING", + createTime: "doc-create-time", + }); + expect(start["status.textQuery"].startTime).toBeDefined(); + expect(start["status.textQuery"].updateTime).toBeDefined(); + + const complete = after.update.mock.calls[1][0]; + expect(complete.result).toEqual({ ids: IDS }); + expect(complete["status.textQuery.state"]).toBe("COMPLETED"); + expect(complete["status.textQuery.updateTime"]).toBeDefined(); + expect(complete["status.textQuery.completeTime"]).toBeDefined(); + }); + + test("keeps an existing status order field across a re-run", async () => { + const { ctx } = makeCtx(); + const after = snapshot({ + query: "second query", + status: { textQuery: { createTime: "first-run-create-time" } }, + }); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "first query" }), after), + ctx + ); + + expect(after.update.mock.calls[0][0]["status.textQuery"].createTime).toBe( + "first-run-create-time" + ); + }); + + test("ignores the result write it makes itself", async () => { + const { ctx } = makeCtx(); + const after = snapshot({ + query: "test query", + result: { ids: IDS }, + ...COMPLETED_STATUS, + }); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "test query" }), after), + ctx + ); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(after.update).not.toHaveBeenCalled(); + }); + + test.each(["PROCESSING", "COMPLETED", "ERROR", "BACKFILLED"])( + "never re-runs a document whose status is %s, even when the query changes", + async (state) => { + const { ctx } = makeCtx(); + const after = snapshot({ + query: "changed query", + status: { textQuery: { state } }, + }); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "original query" }), after), + ctx + ); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(after.update).not.toHaveBeenCalled(); + } + ); + + test("ignores a write that changes neither the query nor the limit", async () => { + const { ctx } = makeCtx(); + const after = snapshot({ query: "test query", unrelated: "b" }); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "test query", unrelated: "a" }), after), + ctx + ); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(after.update).not.toHaveBeenCalled(); + }); + + test("runs when the limit changes on a document with no status", async () => { + const { ctx, chain } = makeCtx(); + const after = snapshot({ query: "test query", limit: "5" }); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "test query", limit: "3" }), after), + ctx + ); + + expect(chain.findNearest).toHaveBeenCalledWith( + config.outputFieldName, + EMBEDDING, + { limit: 5, distanceMeasure: config.distanceMeasure } + ); + }); + + test("applies the document prefilters", async () => { + const { ctx, chain } = makeCtx(); + const after = snapshot({ + query: "test query", + prefilters: [{ field: "category", operator: "==", value: "test" }], + }); + + await handleQueryOnWrite(writeEvent(snapshot(undefined), after), ctx); + + expect(chain.where).toHaveBeenCalledWith("category", "==", "test"); + }); + + test("marks the document ERROR and succeeds when the query fails", async () => { + const { ctx } = makeCtx(); + getSingleEmbedding.mockRejectedValue(new Error("Embedding failed")); + const after = snapshot({ query: "test query" }); + + await expect( + handleQueryOnWrite(writeEvent(snapshot(undefined), after), ctx) + ).resolves.toBeUndefined(); + + expect(after.update).toHaveBeenCalledTimes(2); + const failure = after.update.mock.calls[1][0]; + expect(failure["status.textQuery.state"]).toBe("ERROR"); + expect(failure["status.textQuery.updateTime"]).toBeDefined(); + expect(failure.result).toBeUndefined(); + }); + + test("ignores a document without a string query", async () => { + const { ctx } = makeCtx(); + const after = snapshot({ query: 42 }); + + await handleQueryOnWrite(writeEvent(snapshot(undefined), after), ctx); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(after.update).not.toHaveBeenCalled(); + }); + + test("ignores a delete", async () => { + const { ctx } = makeCtx(); + const after = snapshot(undefined); + + await handleQueryOnWrite( + writeEvent(snapshot({ query: "test query" }), after), + ctx + ); + + expect(getSingleEmbedding).not.toHaveBeenCalled(); + expect(after.update).not.toHaveBeenCalled(); + }); +}); From f7864ad69a29d5af9394fdffc02adc37345d341c Mon Sep 17 00:00:00 2001 From: Corie Watson Date: Wed, 2 Sep 2026 15:03:58 +0100 Subject: [PATCH 2/2] fix(firestore-vector-search): scope the query catch to the query call The extension's processor wrapped only the process function; writeStartEvent and writeCompletionEvent sat outside the try, so a failed status write propagated out of run() and failed the invocation. Wrapping the result write too turned a successful query with a failed result write into a misleading ERROR status write that then failed itself. Narrow the try to performTextQuery. A failed result write now propagates, and the retry skips on the already-written PROCESSING status. Also note in the README that the query status path is literally `status` whatever STATUS_FIELD_NAME is set to, as it was in the extension. --- kits/firestore-vector-search/README.md | 3 ++- kits/firestore-vector-search/src/handlers.ts | 22 ++++++++++++------- .../tests/handlers.test.ts | 14 ++++++++++++ 3 files changed, 30 insertions(+), 9 deletions(-) diff --git a/kits/firestore-vector-search/README.md b/kits/firestore-vector-search/README.md index 92d1d46b7c..92ac59e91f 100644 --- a/kits/firestore-vector-search/README.md +++ b/kits/firestore-vector-search/README.md @@ -253,7 +253,8 @@ updating. The field name is still `STATUS_FIELD_NAME`, defaulting to `status`. This applies to the documents in your indexed collection only. Query documents keep the extension's nested shape, `status.textQuery`, with the same states and -timestamps, so anything waiting on that path still works. +timestamps, so anything waiting on that path still works. As in the extension, +that path is always literally `status`, whatever `STATUS_FIELD_NAME` is set to. ### Editing a document's input re-embeds it diff --git a/kits/firestore-vector-search/src/handlers.ts b/kits/firestore-vector-search/src/handlers.ts index 326a082310..6819e7abc9 100644 --- a/kits/firestore-vector-search/src/handlers.ts +++ b/kits/firestore-vector-search/src/handlers.ts @@ -174,8 +174,12 @@ export async function handleQueryOnWrite( }, }); + // Only the query itself is guarded, as in the extension, where the status + // writes sat outside the processor's try block. A failed status write still + // fails the invocation; the retry sees PROCESSING and skips. + let result: Awaited>; try { - const result = await performTextQuery({ + result = await performTextQuery({ query, limit: data.limit ? parseLimit(data.limit) : ctx.config.defaultQueryLimit, prefilters: (data.prefilters as Prefilter[] | undefined) ?? [], @@ -183,13 +187,6 @@ export async function handleQueryOnWrite( vectorStore: vectorStore(ctx), config: ctx.config, }); - const completeTime = FieldValue.serverTimestamp(); - await after.ref.update({ - ...result, - [`${QUERY_STATUS_PATH}.state`]: "COMPLETED", - [`${QUERY_STATUS_PATH}.updateTime`]: completeTime, - [`${QUERY_STATUS_PATH}.completeTime`]: completeTime, - }); } catch (err) { // The extension's `errorFn` logged and swallowed, so a failed query marks // the document ERROR and the invocation succeeds rather than retrying. @@ -198,7 +195,16 @@ export async function handleQueryOnWrite( [`${QUERY_STATUS_PATH}.state`]: "ERROR", [`${QUERY_STATUS_PATH}.updateTime`]: FieldValue.serverTimestamp(), }); + return; } + + const completeTime = FieldValue.serverTimestamp(); + await after.ref.update({ + ...result, + [`${QUERY_STATUS_PATH}.state`]: "COMPLETED", + [`${QUERY_STATUS_PATH}.updateTime`]: completeTime, + [`${QUERY_STATUS_PATH}.completeTime`]: completeTime, + }); } export async function handleQueryCall( diff --git a/kits/firestore-vector-search/tests/handlers.test.ts b/kits/firestore-vector-search/tests/handlers.test.ts index ce60e34ec4..34c6b20a01 100644 --- a/kits/firestore-vector-search/tests/handlers.test.ts +++ b/kits/firestore-vector-search/tests/handlers.test.ts @@ -399,6 +399,20 @@ describe("handleQueryOnWrite", () => { expect(failure.result).toBeUndefined(); }); + test("propagates a failed result write instead of marking it ERROR", async () => { + const { ctx } = makeCtx(); + const after = snapshot({ query: "test query" }); + after.update + .mockResolvedValueOnce(undefined) + .mockRejectedValueOnce(new Error("Document does not exist")); + + await expect( + handleQueryOnWrite(writeEvent(snapshot(undefined), after), ctx) + ).rejects.toThrow("Document does not exist"); + + expect(after.update).toHaveBeenCalledTimes(2); + }); + test("ignores a document without a string query", async () => { const { ctx } = makeCtx(); const after = snapshot({ query: 42 });