apps/server/src/core/ai-chat/ai-chat.service.setup-abort.spec.ts

@@ -64,6 +64,32 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the
     ],
   };
 
+  // A minimal raw ServerResponse stand-in for the turns that PROCEED past setup
+  // and reach streamText (the deadline + legacy paths). The setup-only abort test
+  // never wires the stream, so it keeps using `{ raw: {} }`.
+  function makeRawRes() {
+    return {
+      raw: {
+        writeHead: jest.fn(function writeHead(this: unknown) {
+          return this;
+        }),
+        write: jest.fn(),
+        once: jest.fn(),
+        flushHeaders: jest.fn(),
+      },
+    };
+  }
+
+  // A fake streamText result: the service only calls consumeStream() and
+  // pipeUIMessageStreamToResponse() on it (both fire-and-forget). Its terminal
+  // callbacks are never invoked, so the run is not finalized through them.
+  function makeStreamResult() {
+    return {
+      consumeStream: jest.fn(),
+      pipeUIMessageStreamToResponse: jest.fn(),
+    };
+  }
+
   beforeEach(() => {
     streamTextMock.mockReset();
     jest.spyOn(Logger.prototype, 'log').mockImplementation(() => undefined as never);
@@ -71,7 +97,10 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the
     jest.spyOn(Logger.prototype, 'error').mockImplementation(() => undefined as never);
   });
 
-  afterEach(() => jest.restoreAllMocks());
+  afterEach(() => {
+    jest.restoreAllMocks();
+    jest.useRealTimers();
+  });
 
   it('stops the hung toolset build, rejects, and settles the run "aborted" — never reaching streamText', async () => {
     const runController = new AbortController();
@@ -118,4 +147,149 @@ describe('AiChatService.stream — abort during external-MCP setup finalizes the
     expect(toolsFor).toHaveBeenCalledTimes(1);
     expect(streamTextMock).not.toHaveBeenCalled();
   });
+
+  // Item 1 — the onLateResolve leg of raceAgainstAbortAndTimeout. When `toolsFor`
+  // loses the race (abort) but RESOLVES LATER with a leased toolset, the setup site
+  // must release that abandoned toolset's leases (call close() on its client
+  // handles) so their lease refcount is not pinned forever by a toolset nobody
+  // consumes. Nothing else exercises this path.
+  it('releases the leases of a toolset that resolves AFTER the race was already lost (onLateResolve)', async () => {
+    const runController = new AbortController();
+    // A controllable build: it hangs until we resolve it by hand, and the run is
+    // stopped mid-build so the race rejects BEFORE the build settles.
+    let resolveTools: (v: unknown) => void = () => undefined;
+    const toolsForPromise = new Promise((resolve) => {
+      resolveTools = resolve;
+    });
+    const toolsFor = jest.fn(() => {
+      setTimeout(() => runController.abort(new Error('user stop')), 0);
+      return toolsForPromise;
+    });
+    const { svc } = makeService({ toolsFor });
+
+    const begin = jest.fn(async () => ({
+      runId: 'run-1',
+      signal: runController.signal,
+    }));
+
+    const promise = svc.stream({
+      user: { id: 'user-1' } as never,
+      workspace: { id: 'ws-1' } as never,
+      sessionId: 'sess-1',
+      body: body as never,
+      res: { raw: {} } as never,
+      signal: new AbortController().signal,
+      model: {} as never,
+      role: null,
+      runHooks: {
+        begin,
+        onAssistantSeeded: jest.fn(),
+        onStep: jest.fn(),
+        onSettled: jest.fn(),
+      } as never,
+    });
+
+    // The race is lost to the abort: the turn rejects with the stop reason.
+    await expect(promise).rejects.toThrow('user stop');
+
+    // NOW the abandoned build resolves late with a leased client. onLateResolve must
+    // release it (call close on the lease handle).
+    const close = jest.fn().mockResolvedValue(undefined);
+    resolveTools({
+      tools: {},
+      clients: [{ close }],
+      outcomes: [],
+      instructions: [],
+    });
+    // Flush the microtasks so work.then -> onLateResolve -> Promise.all(close) runs.
+    await new Promise((r) => setImmediate(r));
+
+    expect(close).toHaveBeenCalledTimes(1);
+  });
+
+  // Item 2 — the PURE DEADLINE branch (MCP_TOOLSET_BUILD_DEADLINE_MS). `toolsFor`
+  // never settles and the run's signal is NOT aborted: the race rejects with a
+  // "setup timed out" error, the catch does NOT re-throw (runId set but signal not
+  // aborted), and the turn PROCEEDS Docmost-only. It must reach streamText (the turn
+  // continues, not wedged) and must NOT be finalized 'aborted'.
+  it('proceeds Docmost-only (reaches streamText) when the build hits the deadline without an abort', async () => {
+    jest.useFakeTimers();
+
+    // The build hangs forever; the run's signal is never aborted.
+    const toolsFor = jest.fn(() => new Promise(() => {}));
+    const { svc } = makeService({ toolsFor });
+
+    streamTextMock.mockReturnValue(makeStreamResult() as never);
+
+    const onSettled = jest.fn();
+    const runSignal = new AbortController().signal; // never aborts
+    const begin = jest.fn(async () => ({ runId: 'run-1', signal: runSignal }));
+
+    const promise = svc.stream({
+      user: { id: 'user-1' } as never,
+      workspace: { id: 'ws-1' } as never,
+      sessionId: 'sess-1',
+      body: body as never,
+      res: makeRawRes() as never,
+      signal: new AbortController().signal,
+      model: {} as never,
+      role: null,
+      runHooks: {
+        begin,
+        onAssistantSeeded: jest.fn(),
+        onStep: jest.fn(),
+        onSettled,
+      } as never,
+    });
+
+    // Advance past the 60s build deadline; advanceTimersByTimeAsync flushes the
+    // promise microtasks between timer fires so the whole setup chain runs.
+    await jest.advanceTimersByTimeAsync(60_001);
+    // The turn does not throw out of setup — it continues to stream.
+    await expect(promise).resolves.toBeUndefined();
+
+    // The turn CONTINUED: streamText was reached (Docmost-only), not wedged.
+    expect(toolsFor).toHaveBeenCalledTimes(1);
+    expect(streamTextMock).toHaveBeenCalledTimes(1);
+    // The run was NOT finalized as aborted (the deadline is not a Stop) — the setup
+    // catch settle path never ran, so onSettled is left to streamText's callbacks.
+    expect(onSettled).not.toHaveBeenCalled();
+  });
+
+  // Item 3 — the LEGACY no-runId path. The catch's re-throw is gated on
+  // `runId && effectiveSignal.aborted`. With NO runId (no runHooks) an abort during
+  // setup must NOT re-throw (runId falsy) — the turn warns + proceeds Docmost-only
+  // and streams, and is never finalized 'aborted' via the re-throw. Locks the
+  // `runId &&` half of the guard.
+  it('does NOT re-throw on a setup abort when there is no runId (legacy path proceeds Docmost-only)', async () => {
+    const socketController = new AbortController();
+    // The build hangs; the SOCKET signal (legacy effectiveSignal) aborts mid-build.
+    const toolsFor = jest.fn(() => {
+      setTimeout(() => socketController.abort(new Error('socket closed')), 0);
+      return new Promise(() => {});
+    });
+    const { svc } = makeService({ toolsFor });
+
+    streamTextMock.mockReturnValue(makeStreamResult() as never);
+
+    // No runHooks => runId undefined, effectiveSignal === the socket signal.
+    const promise = svc.stream({
+      user: { id: 'user-1' } as never,
+      workspace: { id: 'ws-1' } as never,
+      sessionId: 'sess-1',
+      body: body as never,
+      res: makeRawRes() as never,
+      signal: socketController.signal,
+      model: {} as never,
+      role: null,
+    });
+
+    // The turn does NOT reject out of setup (no re-throw on the legacy path).
+    await expect(promise).resolves.toBeUndefined();
+
+    // It proceeded Docmost-only and reached streamText — streamText then observes
+    // the already-aborted socket signal via its own abortSignal.
+    expect(toolsFor).toHaveBeenCalledTimes(1);
+    expect(streamTextMock).toHaveBeenCalledTimes(1);
+  });
 });

apps/server/src/core/ai-chat/ai-chat.service.ts

@@ -63,8 +63,12 @@ const MAX_AGENT_STEPS = 20;
 // per-server connect bound in mcp-clients.service (CONNECT_TIMEOUT_MS): even if
 // that per-server timeout regressed, this outer deadline — together with the run's
 // abort signal — guarantees the setup phase can never wedge a turn at step 0 (the
-// production hang) and the run always finalizes. Generous: the per-server bound
-// should fire first in practice; this only backstops a total build stall.
+// production hang) and the run always finalizes. It stays a TRUE backstop because
+// buildEntry connects to the servers CONCURRENTLY, so the total build time is
+// bounded by the SLOWEST single server (~2×CONNECT_TIMEOUT_MS), not the SUM across
+// them — the per-server bound fires first no matter how many servers are enabled,
+// and this outer deadline only catches a total build stall the per-server bound
+// somehow missed.
 const MCP_TOOLSET_BUILD_DEADLINE_MS = 60_000;
 
 // System-prompt addendum injected ONLY on the final step (see prepareAgentStep).
@@ -791,8 +795,11 @@ export class AiChatService implements OnModuleInit {
         // a hung build would hang the turn at step 0 forever (the production hang),
         // unobservant of an explicit Stop. The deadline is defense-in-depth ABOVE
         // the per-server connect bound in mcp-clients.service. On a LATE resolve
-        // (the race was already lost) close the abandoned toolset's leased clients
-        // so their transports are not leaked.
+        // (the race was already lost) RELEASE the abandoned toolset's leases —
+        // c.close() here is the lease handle, so it decrements the cache entry's
+        // refcount; it does NOT force-close the transports (the cache OWNS the
+        // clients and closes them on TTL/evict). This just prevents the lease
+        // refcount from being pinned >=1 forever by a toolset nobody will consume.
         external = await raceAgainstAbortAndTimeout(
           this.mcpClients.toolsFor(workspace.id),
           effectiveSignal,

apps/server/src/core/ai-chat/external-mcp/mcp-clients.service.ts

@@ -255,54 +255,49 @@ export class McpClientsService {
     const callTimeoutMs = mcpCallTimeoutMs();
     const instructions: McpServerInstruction[] = [];
 
-    for (const server of servers) {
-      // Track the connected client OUTSIDE the try so the catch can close it when
-      // it was obtained but not yet registered in `clients` (e.g. tools() threw or
-      // timed out after a successful connect). `registered` flips only once the
-      // client is owned by `clients` (closed at entry teardown), so the catch never
-      // double-closes a registered client.
+    // Per-server connect+tools result, still tagged with its server so the merge
+    // below can be applied in the SAME order as `servers` (see the parallel note).
+    type PerServerResult =
+      | { ok: true; client: McpClient; guarded: Record<string, Tool> }
+      | { ok: false; reason: string };
+
+    // Connect to (and list tools for) every enabled server CONCURRENTLY, so the
+    // total build time is bounded by the SLOWEST single server (~2×
+    // CONNECT_TIMEOUT_MS: connect + tools), NOT the SUM across servers. The
+    // sequential loop this replaced summed those bounds, so with enough all-timing-
+    // out servers the outer MCP_TOOLSET_BUILD_DEADLINE_MS could fire before the
+    // per-server bounds, dropping ALL external tools and inverting the "per-server
+    // bound is primary, outer is a backstop" invariant. Each server keeps its OWN
+    // try/catch + connectWithTimeout/withTimeout bound + close-on-failure logic; a
+    // failed server is skipped, never fatal. Nothing here mutates the shared
+    // arrays — every result is merged IN SERVER ORDER after Promise.all, so tool-
+    // key precedence/disambiguation, `outcomes`, `instructions` and `clients`
+    // ordering all match the previous sequential behavior exactly.
+    const perServer = async (
+      server: (typeof servers)[number],
+    ): Promise<PerServerResult> => {
+      // Track the connected client so the catch can close it when it was obtained
+      // but tools() then threw/timed out (connectWithTimeout closes its OWN orphan
+      // on a connect timeout, so `client` stays undefined on that path). On success
+      // the client is handed back and registered by the merge below (owned by the
+      // entry, closed at teardown) — so it is never double-closed.
       let client: McpClient | undefined;
-      let registered = false;
       try {
         client = await this.connectWithTimeout(server, CONNECT_TIMEOUT_MS);
         const raw = await withTimeout(client.tools(), CONNECT_TIMEOUT_MS);
-        clients.push(client);
-        registered = true;
         const allow = server.toolAllowlist;
         const picked =
           Array.isArray(allow) && allow.length > 0 ? pick(raw, allow) : raw;
         // Bound each tool's execute with a per-call total-timeout guard before
         // merging, so a single chatty-but-stuck call is aborted after the cap.
         const guarded = wrapToolsWithCallTimeout(picked, callTimeoutMs);
-        // Namespace each tool with the sanitized server name AND disambiguate
-        // against names already merged from earlier servers, so no external
-        // tool is silently overwritten on collision. The returned count drives
-        // whether this server's prompt guidance is included (≥1 tool merged).
-        const merged = this.mergeNamespaced(
-          tools,
-          guarded,
-          server.name,
-          server.id,
-        );
-        outcomes.push({ name: server.name, ok: true });
-        // Include this server's guidance ONLY when it actually contributed at
-        // least one tool the agent can call (allowlist may have filtered all of
-        // them out) AND the admin authored non-blank instructions. The header
-        // prefix is the sanitized server name (= the tool namespace prefix).
-        const guide = server.instructions?.trim();
-        if (merged.count > 0 && guide) {
-          instructions.push({
-            serverName: server.name,
-            toolPrefix: merged.prefix,
-            instructions: guide,
-          });
-        }
+        return { ok: true, client, guarded };
       } catch (err) {
         // A failed server is skipped — the turn proceeds with the rest. If connect
         // returned a live client but a later step (tools()) threw, that client was
         // never registered in `clients`, so close it here or its transport/socket
         // leaks (compounding every 60s cache rebuild during a flaky-server outage).
-        if (client && !registered) {
+        if (client) {
           void client.close().catch(() => undefined);
         }
         // Log a short warning (never the URL/headers) so ops can see degradation,
@@ -311,7 +306,45 @@ export class McpClientsService {
         this.logger.warn(
           `External MCP server "${server.name}" unavailable: ${reason}`,
         );
-        outcomes.push({ name: server.name, ok: false, reason });
+        return { ok: false, reason };
+      }
+    };
+
+    // Promise.all preserves array order regardless of settle order, so `results[i]`
+    // is `servers[i]`'s outcome — the merge below stays deterministic and matches
+    // the old sequential order (later servers still override/disambiguate against
+    // earlier ones on a tool-key clash).
+    const results = await Promise.all(servers.map(perServer));
+    for (let i = 0; i < servers.length; i += 1) {
+      const server = servers[i];
+      const result = results[i];
+      if (result.ok !== true) {
+        outcomes.push({ name: server.name, ok: false, reason: result.reason });
+        continue;
+      }
+      clients.push(result.client);
+      // Namespace each tool with the sanitized server name AND disambiguate
+      // against names already merged from earlier servers, so no external
+      // tool is silently overwritten on collision. The returned count drives
+      // whether this server's prompt guidance is included (≥1 tool merged).
+      const merged = this.mergeNamespaced(
+        tools,
+        result.guarded,
+        server.name,
+        server.id,
+      );
+      outcomes.push({ name: server.name, ok: true });
+      // Include this server's guidance ONLY when it actually contributed at
+      // least one tool the agent can call (allowlist may have filtered all of
+      // them out) AND the admin authored non-blank instructions. The header
+      // prefix is the sanitized server name (= the tool namespace prefix).
+      const guide = server.instructions?.trim();
+      if (merged.count > 0 && guide) {
+        instructions.push({
+          serverName: server.name,
+          toolPrefix: merged.prefix,
+          instructions: guide,
+        });
       }
     }