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 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..097a0be81 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,11 +195,100 @@ describe("makeAiSearchToolDiscoveryProvider", () => { }), ); - it.effect("ignores AI Search chunks whose item key is not tracked locally", () => + 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({ + 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.default", + 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({ + 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 +303,13 @@ 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({ @@ -219,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, + }); }), ); }); @@ -336,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); }, @@ -384,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({ @@ -484,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 5f04b35ae..d90c44301 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 }, @@ -405,14 +421,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, @@ -451,24 +459,7 @@ 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 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 integration = input.namespace?.split(".", 1)[0]; const response = yield* Effect.tryPromise({ try: () => aiSearch.search({ @@ -476,25 +467,55 @@ 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, + ...(integration ? { filters: { integration: { $eq: integration } } } : {}), 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 + // 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); } @@ -575,7 +596,10 @@ export const makeAiSearchToolSearchBackend = ( })), Effect.mapError( (cause) => - new SemanticSearchError({ message: "AI Search query failed.", cause }), + new SemanticSearchError({ + message: "AI Search query failed.", + cause, + }), ), ) : notConfigured(),