From 2fb3cad03a12a73d8bc51bde14e99659fcc1274b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Wed, 5 Aug 2026 16:33:39 +0200 Subject: [PATCH] feat(appkit): cache ai-search queries with OBO-safe keys aiSearch did no caching: defaults had cache disabled and, more importantly, no execute() call passed a cacheKey, so the CacheInterceptor never engaged even when enabled. Wire up query caching mirroring the analytics plugin. - Enable cache in aiSearchDefaults with a short 60s TTL (VS results shift as the index resyncs; cache only long enough to absorb bursts/StrictMode). - Build a per-call cacheKey covering everything that changes results: index, queryText, hashed queryVector, queryType, numResults, post-allowlist columns, filters, reranker. Fold the caller identity in ("global" for SP, user id for OBO) so an OBO query never reads an SP or another user's entry. - Add _resolveQueryParams as the single source of default resolution for both _prepareQuery (payload) and _cacheKeyFor (key), so the key can't drift from what's sent. - Run _prepareQuery inside execute on the programmatic query() path too (the route already did), so a cache hit skips embeddingFn and the VS call. - Leave next-page uncached: a page token is a single-use cursor. - Document caching + per-user isolation in the ai-search plugin docs. Signed-off-by: MarioCadenas --- docs/docs/plugins/ai-search.md | 12 + .../appkit/src/plugins/ai-search/ai-search.ts | 145 +++++++++-- .../appkit/src/plugins/ai-search/defaults.ts | 6 +- .../plugins/ai-search/tests/ai-search.test.ts | 235 +++++++++++++++++- 4 files changed, 366 insertions(+), 32 deletions(-) diff --git a/docs/docs/plugins/ai-search.md b/docs/docs/plugins/ai-search.md index 397e00f3a..43024744f 100644 --- a/docs/docs/plugins/ai-search.md +++ b/docs/docs/plugins/ai-search.md @@ -261,6 +261,18 @@ console.log(result.results); Pass optional overrides as a second argument to `query` to adjust `numResults` or other per-call settings. +## Caching + +Query results are cached with a short TTL (60s) so repeated identical queries — including a component that re-renders or mounts twice — reuse a single Vector Search call instead of hitting the index each time. The next-page route is not cached: a page token is a single-use cursor and already identifies the exact page. + +The cache key covers everything that changes results: the resolved index, `queryText`, `queryVector` (hashed), `queryType`, `numResults`, the resolved `columns`, `filters`, and whether reranking is on. Two queries that differ in any of these are cached separately. + +The TTL is short on purpose: Vector Search results shift as the index resyncs, so results are only reused briefly. + +### Per-user isolation + +For `auth: "on-behalf-of-user"` indexes the caller's identity is part of the cache key, so one user never sees another user's cached results — and an on-behalf-of-user query never reads a service-principal-populated entry. Service-principal indexes share a single cache entry across callers. + ## React hook `useAiSearchQuery` reads the configured indexes from the plugin's client config and posts to the right `/:alias/query` route, so the UI never hardcodes an alias. With one index configured it needs no arguments; pass `{ alias }` to target a specific one. diff --git a/packages/appkit/src/plugins/ai-search/ai-search.ts b/packages/appkit/src/plugins/ai-search/ai-search.ts index 345961672..5f0dfd172 100644 --- a/packages/appkit/src/plugins/ai-search/ai-search.ts +++ b/packages/appkit/src/plugins/ai-search/ai-search.ts @@ -1,11 +1,12 @@ +import { createHash } from "node:crypto"; import type express from "express"; -import type { IAppRouter, PluginExecutionSettings } from "shared"; +import type { CacheConfig, IAppRouter, PluginExecutionSettings } from "shared"; import { AiSearchConnector } from "../../connectors/ai-search/client"; import type { VsQueryParams, VsRawResponse, } from "../../connectors/ai-search/types"; -import { getWorkspaceClient } from "../../context"; +import { getCurrentUserId, getWorkspaceClient } from "../../context"; import { createLogger } from "../../logging/logger"; import { Plugin, toPlugin } from "../../plugin"; import type { PluginManifest } from "../../registry"; @@ -166,22 +167,35 @@ export class AiSearchPlugin extends Plugin { // projection past what the app configured. (query() callers are // trusted and keep the override.) const { columns: _clientColumns, ...safeBody } = body; - const plugin = - indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + const isAsUser = indexConfig.auth === "on-behalf-of-user"; + const plugin = isAsUser ? this.asUser(req) : this; const queryType = safeBody.queryType ?? indexConfig.queryType ?? "hybrid"; + // Fold the caller identity into the cache key so an OBO index never + // serves one user's results to another (or an SP-populated entry). + // SP indexes share a stable "global" entry. `resolveUserId` is only + // reached on the OBO path, where a user header is required. + const executorKey = isAsUser ? this.resolveUserId(req) : "global"; + const cache: CacheConfig = { + ...aiSearchDefaults.cache, + cacheKey: this._cacheKeyFor(safeBody, indexConfig, executorKey), + }; + try { // Prepare inside execute so a self-managed embeddingFn runs in the // same OBO context as the query, not as the service principal. - const result = await plugin.execute(async (signal) => { - const prepared = await this._prepareQuery(safeBody, indexConfig); - return this.connector.query( - getWorkspaceClient(), - { indexName: indexConfig.indexName, ...prepared }, - signal, - ); - }, querySettings); + const result = await plugin.execute( + async (signal) => { + const prepared = await this._prepareQuery(safeBody, indexConfig); + return this.connector.query( + getWorkspaceClient(), + { indexName: indexConfig.indexName, ...prepared }, + signal, + ); + }, + { default: { ...aiSearchDefaults, cache } }, + ); this._sendResult(res, result, queryType); } catch (error) { @@ -230,6 +244,11 @@ export class AiSearchPlugin extends Plugin { const plugin = indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + // Not cached: a page token is a single-use cursor into one query's + // result stream, so a repeat request with the same token is rare and + // the token itself already encodes the exact page. `querySettings` + // carries `cache.enabled` but no `cacheKey`, so the interceptor stays + // off (see plugin `_buildInterceptors`). const result = await plugin.execute( async (signal) => this.connector.queryNextPage( @@ -301,16 +320,31 @@ export class AiSearchPlugin extends Plugin { throw new Error(`No index configured with alias "${alias}"`); } - const prepared = await this._prepareQuery(request, indexConfig); + // Resolved up front only for the response's queryType — `_prepareQuery` + // (and any `embeddingFn`) now runs inside `execute`, so a cache hit skips + // both the embedding and the VS call, mirroring the HTTP query route. + const { queryType } = this._resolveQueryParams(request, indexConfig); + + // `query()` runs inside the caller's context — under `asUser(req)` the + // proxy wraps it in `runInUserContext`, so `getCurrentUserId()` is the + // user's id here and folds into the key (SP callers get the service id). + // This matches the interceptor's own `context.userKey`, so an OBO caller + // never shares an SP or another user's cache entry. + const cache: CacheConfig = { + ...aiSearchDefaults.cache, + cacheKey: this._cacheKeyFor(request, indexConfig, getCurrentUserId()), + }; const result = await this.execute( - async (signal) => - this.connector.query( + async (signal) => { + const prepared = await this._prepareQuery(request, indexConfig); + return this.connector.query( getWorkspaceClient(), { indexName: indexConfig.indexName, ...prepared }, signal, - ), - querySettings, + ); + }, + { default: { ...aiSearchDefaults, cache } }, ); if (!result.ok) { @@ -319,7 +353,7 @@ export class AiSearchPlugin extends Plugin { ); } - return this._parseResponse(result.data, prepared.queryType); + return this._parseResponse(result.data, queryType); } async shutdown(): Promise { @@ -373,11 +407,38 @@ export class AiSearchPlugin extends Plugin { res.json(this._parseResponse(result.data, queryType)); } + /** + * Resolve request-vs-index defaults for the fields that determine a query's + * result: `queryType`, `columns` (post-allowlist), `numResults`, and the + * derived `reranker`. Single source of truth for both `_prepareQuery` (what + * gets sent to VS) and `_cacheKeyFor` (what the result is keyed on), so the + * key can never silently drift from the payload. + */ + private _resolveQueryParams( + request: SearchRequest, + indexConfig: IndexConfig, + ): { + queryType: "ann" | "hybrid" | "full_text"; + columns: string[]; + numResults: number; + reranker: { columnsToRerank: string[] } | undefined; + } { + const queryType = request.queryType ?? indexConfig.queryType ?? "hybrid"; + const columns = request.columns ?? indexConfig.columns ?? []; + return { + queryType, + columns, + numResults: request.numResults ?? indexConfig.numResults ?? 20, + reranker: this._resolveReranker(request.reranker, indexConfig, columns), + }; + } + private async _prepareQuery( request: SearchRequest, indexConfig: IndexConfig, ): Promise> { - const queryType = request.queryType ?? indexConfig.queryType ?? "hybrid"; + const { queryType, columns, numResults, reranker } = + this._resolveQueryParams(request, indexConfig); let queryText = request.queryText; let queryVector = request.queryVector; @@ -398,18 +459,58 @@ export class AiSearchPlugin extends Plugin { } } - const columns = request.columns ?? indexConfig.columns ?? []; return { queryText, queryVector, queryType, columns, - numResults: request.numResults ?? indexConfig.numResults ?? 20, + numResults, filters: request.filters, - reranker: this._resolveReranker(request.reranker, indexConfig, columns), + reranker, }; } + /** + * Build the cache key for a query. Every input that changes the VS result + * is included, resolved through the same `_resolveQueryParams` that builds + * the payload so the key always agrees with what's actually sent (e.g. the + * post-allowlist `columns`, not the raw request). `executorKey` folds in the + * caller identity — "global" for the service principal, the user id for OBO + * — so an OBO caller never shares an SP or another user's entry. + * + * `queryVector` is hashed rather than embedded inline (vectors are large). + * We key on the request's inputs, not the derived embedding: an + * `embeddingFn` is a deterministic function of `queryText`, which is already + * in the key, and it hasn't run yet at key-build time (it runs inside + * `execute`, in the OBO context). + */ + private _cacheKeyFor( + request: SearchRequest, + indexConfig: IndexConfig & { indexName: string }, + executorKey: string, + ): CacheConfig["cacheKey"] { + const { queryType, columns, numResults, reranker } = + this._resolveQueryParams(request, indexConfig); + return [ + "ai-search:query", + indexConfig.indexName, + request.queryText ?? "", + request.queryVector ? this._hashVector(request.queryVector) : "", + queryType, + numResults, + JSON.stringify(columns), + JSON.stringify(request.filters ?? null), + // Reranker changes result ordering; keep it in the key. Stringified + // because `cacheKey` entries are string | number | object, not boolean. + String(!!reranker), + executorKey, + ]; + } + + private _hashVector(vector: number[]): string { + return createHash("sha256").update(JSON.stringify(vector)).digest("hex"); + } + private _resolveReranker( requestReranker: boolean | undefined, indexConfig: IndexConfig, diff --git a/packages/appkit/src/plugins/ai-search/defaults.ts b/packages/appkit/src/plugins/ai-search/defaults.ts index d927c4490..2af9e0768 100644 --- a/packages/appkit/src/plugins/ai-search/defaults.ts +++ b/packages/appkit/src/plugins/ai-search/defaults.ts @@ -1,7 +1,11 @@ import type { PluginExecuteConfig } from "shared"; export const aiSearchDefaults: PluginExecuteConfig = { - cache: { enabled: false }, + // Short TTL: Vector Search results shift as the index resyncs, so cache only + // long enough to absorb bursts (repeated queries, React StrictMode double + // mounts) without serving stale results for long. The interceptor still needs + // a per-call `cacheKey` — set in the query paths — before caching engages. + cache: { enabled: true, ttl: 60 }, retry: { enabled: true, initialDelay: 1000, attempts: 3 }, timeout: 30_000, }; diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index 3a8efc0ab..d2768cdd3 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -9,6 +9,17 @@ import { Context } from "../../../workspace-client"; vi.mock("../../../context", () => ({ getWorkspaceClient: vi.fn(() => mockWorkspaceClient), getCurrentUserId: vi.fn(() => "test-user"), + // Enough of the OBO plumbing for asUser() to run its real (non-dev) path: + // createUserContext returns a stub context and runInUserContext just invokes + // the fn. getCurrentUserId stays constant, so per-user cache scoping is + // driven entirely by the executorKey folded into the cacheKey array. + runInUserContext: (_ctx: unknown, fn: () => T): T => fn(), + ServiceContext: { + createUserContext: (_token: string, userId: string) => ({ + userId, + isUserContext: true, + }), + }, })); vi.mock("../../../logging/logger", () => ({ @@ -53,17 +64,39 @@ vi.mock("../../../telemetry", () => ({ normalizeTelemetryOptions: () => ({ traces: false, metrics: false }), })); -vi.mock("../../../cache", () => ({ - CacheManager: { - getInstanceSync: () => ({ - get: vi.fn(), - set: vi.fn(), - delete: vi.fn(), - generateKey: vi.fn(() => "test-key"), - }), - }, +// Faithful in-memory cache: keyed by [userKey, ...cacheKey] like the real +// CacheManager.generateKey, so tests exercise real key composition (hit on an +// identical key, miss on any differing part). Never stores rejections. +const { mockCacheStore } = vi.hoisted(() => ({ + mockCacheStore: new Map(), })); +vi.mock("../../../cache", () => { + const keyOf = (parts: unknown[], userKey: string) => + JSON.stringify([userKey, ...parts]); + return { + CacheManager: { + getInstanceSync: () => ({ + get: vi.fn(), + set: vi.fn(), + delete: vi.fn(), + generateKey: keyOf, + getOrExecute: async ( + key: unknown[], + fn: (signal?: AbortSignal) => Promise, + userKey: string, + ) => { + const k = keyOf(key, userKey); + if (mockCacheStore.has(k)) return mockCacheStore.get(k); + const result = await fn(); + mockCacheStore.set(k, result); + return result; + }, + }), + }, + }; +}); + vi.mock("../../../app", () => ({ AppManager: vi.fn().mockImplementation(() => ({})), })); @@ -108,6 +141,7 @@ describe("AiSearchPlugin", () => { beforeEach(() => { mockRequest.mockClear(); mockRequest.mockResolvedValue(validVsResponse); + mockCacheStore.clear(); }); describe("setup()", () => { @@ -667,6 +701,189 @@ describe("AiSearchPlugin", () => { }); }); + describe("caching", () => { + const makePlugin = () => + new AiSearchPlugin({ + indexes: { + products: { + indexName: "cat.sch.products", + columns: ["id", "title"], + queryType: "hybrid", + numResults: 10, + }, + }, + }); + + it("serves an identical query from cache (connector called once)", async () => { + const plugin = makePlugin(); + await plugin.setup(); + + await plugin.query("products", { queryText: "machine learning" }); + await plugin.query("products", { queryText: "machine learning" }); + + expect(mockRequest).toHaveBeenCalledTimes(1); + }); + + it.each([ + ["queryText", { queryText: "different" }], + ["numResults", { queryText: "q", numResults: 5 }], + ["queryType", { queryText: "q", queryType: "ann" as const }], + ["columns", { queryText: "q", columns: ["id"] }], + ["filters", { queryText: "q", filters: { category: ["books"] } }], + ["reranker", { queryText: "q", reranker: true }], + ])( + "does not share cache entries when %s differs (2 connector calls)", + async (_field, second) => { + const plugin = makePlugin(); + await plugin.setup(); + + await plugin.query("products", { queryText: "q" }); + await plugin.query("products", second); + + expect(mockRequest).toHaveBeenCalledTimes(2); + }, + ); + + it("keys managed-embedding queries by queryText, skipping embedding on a route cache hit", async () => { + // The key is built from the request's queryText, not the derived vector, + // so identical text hits cache. On the route path _prepareQuery (and thus + // embeddingFn) runs inside execute, so a hit skips both embedding and the + // connector call. + const embeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new AiSearchPlugin({ + indexes: { + docs: { + indexName: "cat.sch.docs", + columns: ["id", "title"], + queryType: "ann", + embeddingFn, + }, + }, + }); + await plugin.setup(); + + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const handler = getHandler("POST", "/:alias/query"); + const run = () => + handler( + createMockRequest({ + params: { alias: "docs" }, + body: { queryText: "same" }, + }), + createMockResponse(), + ); + + await run(); + await run(); + + expect(mockRequest).toHaveBeenCalledTimes(1); + expect(embeddingFn).toHaveBeenCalledTimes(1); + }); + + it("skips embedding on a programmatic query() cache hit too", async () => { + // Programmatic query() also runs _prepareQuery inside execute, so a hit + // on the second identical call reuses the cached result without calling + // embeddingFn or the connector again. + const embeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new AiSearchPlugin({ + indexes: { + docs: { + indexName: "cat.sch.docs", + columns: ["id", "title"], + queryType: "ann", + embeddingFn, + }, + }, + }); + await plugin.setup(); + + await plugin.query("docs", { queryText: "same" }); + await plugin.query("docs", { queryText: "same" }); + + expect(mockRequest).toHaveBeenCalledTimes(1); + expect(embeddingFn).toHaveBeenCalledTimes(1); + }); + + it("does not share cache across OBO users (per-user cache key)", async () => { + const plugin = new AiSearchPlugin({ + indexes: { + docs: { + indexName: "cat.sch.docs", + columns: ["id", "title"], + auth: "on-behalf-of-user", + }, + }, + }); + await plugin.setup(); + + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const handler = getHandler("POST", "/:alias/query"); + + const runAs = async (user: string) => { + const res = createMockResponse(); + await handler( + createMockRequest({ + params: { alias: "docs" }, + body: { queryText: "shared question" }, + headers: { + "x-forwarded-user": user, + "x-forwarded-access-token": "tok", + }, + }), + res, + ); + return res; + }; + + const resA = await runAs("alice"); + const resB = await runAs("bob"); + + // Same query text, different users → distinct keys → both hit the + // connector. A shared entry would collapse this to one call and leak + // alice's results to bob. + expect(mockRequest).toHaveBeenCalledTimes(2); + expect(resA.json).toHaveBeenCalled(); + expect(resB.json).toHaveBeenCalled(); + }); + + it("re-serves the same OBO user from cache (connector called once)", async () => { + const plugin = new AiSearchPlugin({ + indexes: { + docs: { + indexName: "cat.sch.docs", + columns: ["id", "title"], + auth: "on-behalf-of-user", + }, + }, + }); + await plugin.setup(); + + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + const handler = getHandler("POST", "/:alias/query"); + + const runAsAlice = () => + handler( + createMockRequest({ + params: { alias: "docs" }, + body: { queryText: "shared question" }, + headers: { + "x-forwarded-user": "alice", + "x-forwarded-access-token": "tok", + }, + }), + createMockResponse(), + ); + + await runAsAlice(); + await runAsAlice(); + + expect(mockRequest).toHaveBeenCalledTimes(1); + }); + }); + describe("injectRoutes", () => { const makePlugin = () => new AiSearchPlugin({