From af756e54eac1630666ed13efdca929bc7dc5035e Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Sun, 2 Aug 2026 20:34:51 +0530 Subject: [PATCH 1/4] fix(semantic-search): reconcile provider item paths --- .../semantic-search/src/sdk/ai-search.test.ts | 73 +++++++++++++++++-- .../semantic-search/src/sdk/ai-search.ts | 64 ++++++++-------- 2 files changed, 96 insertions(+), 41 deletions(-) diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts index a2eb6493f..fa3f1a79e 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts @@ -183,11 +183,71 @@ describe("makeAiSearchToolDiscoveryProvider", () => { }), ); - it.effect("ignores AI Search chunks whose item key is not tracked locally", () => + it.effect("reconciles a provider-rewritten item key by canonical path", () => + Effect.gen(function* () { + const provider = makeAiSearchToolDiscoveryProvider({ + aiSearch: { + ...makeAiSearch(), + search: async () => ({ + search_query: "create repo", + chunks: [ + { + id: "chunk-1", + type: "text", + score: 0.7, + text: "create a repository", + item: { + key: "provider-rewritten-key.md", + metadata: { + path: githubRow.key, + name: githubRow.data.name, + description: githubRow.data.description, + integration: githubRow.data.integration, + }, + }, + }, + ], + }), + }, + items: makeItemsCollection({ + getMany: ({ keys }) => + Effect.succeed( + new Map( + keys.flatMap((key) => (key === githubRow.key ? [[key, githubRow] as const] : [])), + ), + ), + }), + }); + + const page = yield* provider!.searchTools({ + executor: undefined as never, + query: "create repo", + limit: 10, + offset: 0, + }); + + expect(page.items).toMatchObject([ + { + path: githubRow.key, + name: githubRow.data.name, + score: 0.7, + }, + ]); + }), + ); + + it.effect("ignores AI Search chunks whose paths are not current locally", () => Effect.gen(function* () { const provider = makeAiSearchToolDiscoveryProvider({ aiSearch: makeAiSearch(), - items: makeItemsCollection({ list: () => Effect.succeed([githubRow]) }), + items: makeItemsCollection({ + getMany: ({ keys }) => + Effect.succeed( + new Map( + keys.flatMap((key) => (key === githubRow.key ? [[key, githubRow] as const] : [])), + ), + ), + }), }); const page = yield* provider!.searchTools({ @@ -202,14 +262,11 @@ describe("makeAiSearchToolDiscoveryProvider", () => { }), ); - it.effect("returns an empty page without querying AI Search when local rows are empty", () => + it.effect("returns an empty page when no returned paths are current locally", () => Effect.gen(function* () { const provider = makeAiSearchToolDiscoveryProvider({ - aiSearch: { - ...makeAiSearch(), - search: () => expect.unreachable("AI Search should not be queried"), - }, - items: makeItemsCollection({ list: () => Effect.succeed([]) }), + aiSearch: makeAiSearch(), + items: makeItemsCollection({ getMany: () => Effect.succeed(new Map()) }), }); const page = yield* provider!.searchTools({ diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.ts b/packages/plugins/semantic-search/src/sdk/ai-search.ts index 5f04b35ae..bf665741c 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.ts @@ -405,14 +405,6 @@ export const statusAiSearch = (input: { const matchesNamespace = (path: string, namespace: string | undefined): boolean => !namespace || path === namespace || path.startsWith(`${namespace}.`); -const rowToResult = (row: AiSearchItemRow, score: number): ToolDiscoveryResult => ({ - path: row.path, - name: row.name, - description: row.description, - integration: row.integration, - score, -}); - const getStringMetadata = ( metadata: Readonly> | undefined, key: string, @@ -452,23 +444,6 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { return { items: [], total: 0, hasMore: false, nextOffset: null }; } const limit = Math.min(50, Math.max(1, input.limit + input.offset)); - const hasLocalRows = deps.items !== undefined; - const rowsByKey = - deps.items === undefined - ? undefined - : yield* deps.items.list().pipe( - Effect.map((rows) => new Map(rows.map((row) => [row.data.key, row.data]))), - Effect.mapError( - (cause) => - new ExecutionToolError({ - message: "AI Search tool search failed.", - cause, - }), - ), - ); - if (hasLocalRows && rowsByKey?.size === 0) { - return { items: [], total: 0, hasMore: false, nextOffset: null }; - } const response = yield* Effect.tryPromise({ try: () => aiSearch.search({ @@ -486,15 +461,38 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { new ExecutionToolError({ message: "AI Search tool search failed.", cause }), }); + // AI Search carries the canonical tool metadata on every chunk. Validate its + // path against the current catalog, rather than reconciling its opaque item key + // with the local upload ledger. That key is provider-owned and may be rewritten + // during an upload, while the catalog path is the stable identity clients use. + const chunkResults = (response.chunks ?? []).flatMap((chunk) => { + const result = chunkToResult(chunk); + return result === null ? [] : [result]; + }); + const visiblePaths = + deps.items === undefined + ? undefined + : new Set( + (yield* deps.items + .getMany({ keys: chunkResults.map((result) => result.path) }) + .pipe( + Effect.mapError( + (cause) => + new ExecutionToolError({ + message: "AI Search tool search failed.", + cause, + }), + ), + )).keys(), + ); + const bestByPath = new Map(); - for (const chunk of response.chunks ?? []) { - const row = chunk.item?.key ? rowsByKey?.get(chunk.item.key) : undefined; - const result = row - ? rowToResult(row, chunk.score) - : hasLocalRows - ? null - : chunkToResult(chunk); - if (!result || !matchesNamespace(result.path, input.namespace)) continue; + for (const result of chunkResults) { + if ( + (visiblePaths !== undefined && !visiblePaths.has(result.path)) || + !matchesNamespace(result.path, input.namespace) + ) + continue; const previous = bestByPath.get(result.path); if (!previous || result.score > previous.score) bestByPath.set(result.path, result); } From a26043c9a33a697abddeacd77b0b54827c4b54ed Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Mon, 10 Aug 2026 19:59:23 +0530 Subject: [PATCH 2/4] fix(cloudflare): avoid restoring runtime during owner validation Keep Durable Object owner checks storage-only so PartyServer remains the sole runtime initialization path and concurrent fetches cannot deadlock startup. --- .../mcp/agent-session-durable-object.test.ts | 32 ++++++++----------- .../src/mcp/agent-session-durable-object.ts | 20 ++++++------ 2 files changed, 24 insertions(+), 28 deletions(-) diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts index fe70e4cb2..3661bd74d 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts @@ -307,32 +307,30 @@ describe("McpAgentSessionDOBase apps capability persistence", () => { }); describe("McpAgentSessionDOBase transport restore", () => { - it("restores a same-session request after idle disposal leaves a stale server transport", async () => { + it("validates a same-session request after idle disposal without restoring the runtime", async () => { const session = await makeHarnessSession(); + let onStartCalls = 0; + session.runMcpAgentOnStart = async () => { + onStartCalls += 1; + }; await session.alarm(); await expect( session.validateMcpSessionOwner({ accountId: "user-1", organizationId: "org-1" }), ).resolves.toBe("ok"); + expect(onStartCalls).toBe(0); + expect(session.initialized).toBe(false); + expect(session.engine).toBeNull(); + expect(session.server).toBeUndefined(); }); - it("single-flights concurrent same-session restore after idle disposal", async () => { + it("does not start runtime for concurrent owner validations after idle disposal", async () => { const session = await makeHarnessSession(); - const firstRestoreEntered = makeDeferred(); - const finishRestore = makeDeferred(); let onStartCalls = 0; - let restoredServer: McpServer | undefined; session.runMcpAgentOnStart = async () => { onStartCalls += 1; - const restored = session.server ?? makeServer(); - restoredServer ??= restored; - session.server = restored; - firstRestoreEntered.resolve(); - await finishRestore.promise; - await restored.connect(new RestoredTransport()); - session.initialized = true; }; await session.alarm(); @@ -346,16 +344,12 @@ describe("McpAgentSessionDOBase transport restore", () => { organizationId: "org-1", }); - await firstRestoreEntered.promise; - await Promise.resolve(); - finishRestore.resolve(); - await expect(Promise.all([first, second])).resolves.toEqual(["ok", "ok"]); - expect(onStartCalls).toBe(1); - expect(session.server).toBe(restoredServer); + expect(onStartCalls).toBe(0); + expect(session.initialized).toBe(false); }); - it("single-flights SDK onStart callers with same-session restore", async () => { + it("does not race SDK onStart with owner validation", async () => { const session = await makeHarnessSession(); const firstStartEntered = makeDeferred(); const finishStart = makeDeferred(); diff --git a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts index e29c86628..cf8124a3d 100644 --- a/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts +++ b/packages/hosts/cloudflare/src/mcp/agent-session-durable-object.ts @@ -763,19 +763,21 @@ export abstract class McpAgentSessionDOBase< if (destroyPending === true) return "terminated" as const; const sessionMeta = yield* self.loadSessionMeta(); if (!sessionMeta) return "not_found" as const; - if (self.initialized) { + const owner = + identity.accountId === sessionMeta.userId && + identity.organizationId === sessionMeta.organizationId + ? ("ok" as const) + : ("forbidden" as const); + if (owner === "ok" && self.initialized) { yield* Effect.promise(() => self.markActivity()).pipe( Effect.withSpan("McpSessionDO.markActivity"), ); - } else { - yield* Effect.promise(() => self.onStart()).pipe( - Effect.withSpan("McpSessionDO.restore_transport_runtime"), - ); } - return identity.accountId === sessionMeta.userId && - identity.organizationId === sessionMeta.organizationId - ? ("ok" as const) - : ("forbidden" as const); + // Owner validation is intentionally storage-only. A cold session is + // initialized by PartyServer's fetch path below; starting it here + // would run outside PartyServer's single blockConcurrencyWhile lock, + // so a concurrent fetch can wait on the same DO startup indefinitely. + return owner; }).pipe( Effect.withSpan("McpSessionDO.validateMcpSessionOwner"), // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: DO RPC exposes Promise results From 9fc72614a599bcc410ea1cc99728c6f476deb26b Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Mon, 10 Aug 2026 19:59:37 +0530 Subject: [PATCH 3/4] fix(semantic-search): harden AI Search indexing and retrieval Re-upload unexpected remote item states, preserve queued/running/completed items, and filter scoped searches inside AI Search before tool-level deduplication and paging. --- .../semantic-search/src/sdk/ai-search.test.ts | 76 +++++++++++++++++-- .../semantic-search/src/sdk/ai-search.ts | 49 +++++++++--- 2 files changed, 106 insertions(+), 19 deletions(-) diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts index fa3f1a79e..da8f23e97 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts @@ -64,15 +64,27 @@ const makeItemsCollection = (overrides: Partial): ItemsCollecti const makeAiSearchItems = () => ({ - upload: async (name) => ({ id: `item:${name}`, key: name, status: "queued" }), + upload: async (name) => ({ + id: `item:${name}`, + key: name, + status: "queued", + }), list: async () => ({ result: [], result_info: { count: 0, total_count: 0, page: 1, per_page: 50 }, }), delete: async () => {}, - uploadAndPoll: async (name) => ({ id: `item:${name}`, key: name, status: "queued" }), + uploadAndPoll: async (name) => ({ + id: `item:${name}`, + key: name, + status: "queued", + }), get: (itemId) => ({ - info: async () => ({ id: itemId, key: itemId.replace(/^item:/, ""), status: "queued" }), + info: async () => ({ + id: itemId, + key: itemId.replace(/^item:/, ""), + status: "queued", + }), download: async () => expect.unreachable("Unexpected AI Search item download"), }), }) satisfies Pick["items"]; @@ -183,6 +195,35 @@ describe("makeAiSearchToolDiscoveryProvider", () => { }), ); + it.effect("pushes an integration namespace into AI Search retrieval", () => + Effect.gen(function* () { + let request: Parameters[0] | undefined; + const provider = makeAiSearchToolDiscoveryProvider({ + aiSearch: { + ...makeAiSearch(), + search: async (input) => { + request = input; + return makeAiSearch().search(input); + }, + }, + items: undefined, + }); + + yield* provider!.searchTools({ + executor: undefined as never, + query: "authenticated user", + namespace: "github_api", + limit: 5, + offset: 0, + }); + + expect(request?.ai_search_options?.retrieval).toMatchObject({ + max_num_results: 50, + filters: { integration: { $eq: "github_api" } }, + }); + }), + ); + it.effect("reconciles a provider-rewritten item key by canonical path", () => Effect.gen(function* () { const provider = makeAiSearchToolDiscoveryProvider({ @@ -266,7 +307,9 @@ describe("makeAiSearchToolDiscoveryProvider", () => { Effect.gen(function* () { const provider = makeAiSearchToolDiscoveryProvider({ aiSearch: makeAiSearch(), - items: makeItemsCollection({ getMany: () => Effect.succeed(new Map()) }), + items: makeItemsCollection({ + getMany: () => Effect.succeed(new Map()), + }), }); const page = yield* provider!.searchTools({ @@ -276,7 +319,12 @@ describe("makeAiSearchToolDiscoveryProvider", () => { offset: 0, }); - expect(page).toMatchObject({ items: [], total: 0, hasMore: false, nextOffset: null }); + expect(page).toMatchObject({ + items: [], + total: 0, + hasMore: false, + nextOffset: null, + }); }), ); }); @@ -393,7 +441,11 @@ describe("reindexAiSearch", () => { ...makeAiSearch(), items: { ...makeAiSearchItems(), - upload: async (name) => ({ id: `new:${name}`, key: name, status: "queued" }), + upload: async (name) => ({ + id: `new:${name}`, + key: name, + status: "queued", + }), delete: async (id) => { deleted.push(id); }, @@ -441,7 +493,11 @@ describe("reindexAiSearch", () => { ...makeAiSearch(), items: { ...makeAiSearchItems(), - upload: async (name) => ({ id: `item:${name}`, key: name, status: "queued" }), + upload: async (name) => ({ + id: `item:${name}`, + key: name, + status: "queued", + }), }, }, items: makeItemsCollection({ @@ -541,7 +597,11 @@ describe("reindexAiSearch", () => { ...makeAiSearch(), items: { ...makeAiSearchItems(), - upload: async (name) => ({ id: `item:${name}`, key: name, status: "queued" }), + upload: async (name) => ({ + id: `item:${name}`, + key: name, + status: "queued", + }), }, }, items: makeItemsCollection({ diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.ts b/packages/plugins/semantic-search/src/sdk/ai-search.ts index bf665741c..4ab9c1cc9 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.ts @@ -50,10 +50,15 @@ const AI_SEARCH_UPLOAD_BATCH_SIZE = 25; const nowIso = (): string => new Date().toISOString(); +const isReusableRemoteStatus = (status: string | undefined): boolean => + status === undefined || status === "queued" || status === "running" || status === "completed"; + const toStatus = (status: string | undefined): AiSearchItemStatus => status === "queued" || status === "running" || status === "completed" || status === "error" ? status - : "queued"; + : status === undefined + ? "queued" + : "error"; const toItemName = (document: ToolSearchDocument): string => `tool-${cyrb53(`${document.path}\u0000${document.fingerprint}`).toString(36)}.md`; @@ -105,7 +110,10 @@ const deleteItem = ( Effect.tryPromise({ try: () => aiSearch.items.delete(itemId), catch: (cause) => - new SemanticSearchError({ message: `Failed to delete AI Search item "${itemId}".`, cause }), + new SemanticSearchError({ + message: `Failed to delete AI Search item "${itemId}".`, + cause, + }), }).pipe(Effect.asVoid); const deleteItemBestEffort = ( @@ -120,7 +128,10 @@ const getAiSearchItem = ( Effect.tryPromise({ try: () => aiSearch.items.get(itemId).info(), catch: (cause) => - new SemanticSearchError({ message: `Failed to get AI Search item "${itemId}".`, cause }), + new SemanticSearchError({ + message: `Failed to get AI Search item "${itemId}".`, + cause, + }), }); const toIndexedItemRow = ( @@ -156,7 +167,7 @@ const uploadDocument = ( ): Effect.Effect => Effect.gen(function* () { const itemName = toItemName(document); - if (remote !== undefined && remote.status !== "error") { + if (remote !== undefined && isReusableRemoteStatus(remote.status)) { return { deleteOnStorageFailure: false, uploadedItemId: remote.id, @@ -202,7 +213,9 @@ export const reindexAiSearchBatch = (input: { const aiSearch = input.aiSearch; return Effect.gen(function* () { const batch = normalizeBatchInput(input); - const manifests = yield* listToolManifests(input.executor, { maxTools: batch.maxTools }); + const manifests = yield* listToolManifests(input.executor, { + maxTools: batch.maxTools, + }); const page = manifests.slice(batch.offset, batch.offset + batch.pageSize); const nextOffset = batch.offset + page.length < manifests.length ? batch.offset + page.length : null; @@ -245,7 +258,7 @@ export const reindexAiSearchBatch = (input: { if ( previous?.fingerprint === fingerprint && remote !== undefined && - remote.status !== "error" + isReusableRemoteStatus(remote.status) ) { skipped += 1; continue; @@ -383,7 +396,10 @@ export const statusAiSearch = (input: { Effect.tryPromise({ try: () => input.aiSearch.stats(), catch: (cause) => - new SemanticSearchError({ message: "Failed to read AI Search status.", cause }), + new SemanticSearchError({ + message: "Failed to read AI Search status.", + cause, + }), }), ] as const, { concurrency: 2 }, @@ -443,7 +459,6 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { if (!query) { return { items: [], total: 0, hasMore: false, nextOffset: null }; } - const limit = Math.min(50, Math.max(1, input.limit + input.offset)); const response = yield* Effect.tryPromise({ try: () => aiSearch.search({ @@ -451,14 +466,23 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { ai_search_options: { retrieval: { retrieval_type: "hybrid", - max_num_results: limit, + // Retrieve a broad candidate set before deduplication and paging. Asking + // AI Search for only the caller's page size makes plausible tools vanish + // when several chunks belong to one tool or beat the desired integration. + max_num_results: 50, + ...(input.namespace + ? { filters: { integration: { $eq: input.namespace } } } + : {}), return_on_failure: true, }, reranking: { enabled: true }, }, }), catch: (cause) => - new ExecutionToolError({ message: "AI Search tool search failed.", cause }), + new ExecutionToolError({ + message: "AI Search tool search failed.", + cause, + }), }); // AI Search carries the canonical tool metadata on every chunk. Validate its @@ -573,7 +597,10 @@ export const makeAiSearchToolSearchBackend = ( })), Effect.mapError( (cause) => - new SemanticSearchError({ message: "AI Search query failed.", cause }), + new SemanticSearchError({ + message: "AI Search query failed.", + cause, + }), ), ) : notConfigured(), From cac6dc2a199deb6d920e17f6eb3d1726d7c0797e Mon Sep 17 00:00:00 2001 From: Saatvik Arya Date: Mon, 10 Aug 2026 20:13:16 +0530 Subject: [PATCH 4/4] fix(semantic-search): scope provider filters by integration (greptile) --- packages/plugins/semantic-search/src/sdk/ai-search.test.ts | 4 ++-- packages/plugins/semantic-search/src/sdk/ai-search.ts | 5 ++--- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts index da8f23e97..097a0be81 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.test.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.test.ts @@ -195,7 +195,7 @@ describe("makeAiSearchToolDiscoveryProvider", () => { }), ); - it.effect("pushes an integration namespace into AI Search retrieval", () => + it.effect("pushes the integration segment of a path namespace into AI Search retrieval", () => Effect.gen(function* () { let request: Parameters[0] | undefined; const provider = makeAiSearchToolDiscoveryProvider({ @@ -212,7 +212,7 @@ describe("makeAiSearchToolDiscoveryProvider", () => { yield* provider!.searchTools({ executor: undefined as never, query: "authenticated user", - namespace: "github_api", + namespace: "github_api.default", limit: 5, offset: 0, }); diff --git a/packages/plugins/semantic-search/src/sdk/ai-search.ts b/packages/plugins/semantic-search/src/sdk/ai-search.ts index 4ab9c1cc9..d90c44301 100644 --- a/packages/plugins/semantic-search/src/sdk/ai-search.ts +++ b/packages/plugins/semantic-search/src/sdk/ai-search.ts @@ -459,6 +459,7 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { if (!query) { return { items: [], total: 0, hasMore: false, nextOffset: null }; } + const integration = input.namespace?.split(".", 1)[0]; const response = yield* Effect.tryPromise({ try: () => aiSearch.search({ @@ -470,9 +471,7 @@ export const makeAiSearchToolDiscoveryProvider = (deps: { // AI Search for only the caller's page size makes plausible tools vanish // when several chunks belong to one tool or beat the desired integration. max_num_results: 50, - ...(input.namespace - ? { filters: { integration: { $eq: input.namespace } } } - : {}), + ...(integration ? { filters: { integration: { $eq: integration } } } : {}), return_on_failure: true, }, reranking: { enabled: true },