fix(selfhost): reconcile MCP runtime and AI Search

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();

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

packages/plugins/semantic-search/src/sdk/ai-search.test.ts

@@ -64,15 +64,27 @@ const makeItemsCollection = (overrides: Partial<ItemsCollection>): 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<AiSearchInstance, "items">["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<AiSearchInstance["search"]>[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({