|
129 | 129 | (future (f jsonrpc-notification))))) |
130 | 130 |
|
131 | 131 | (defn ^:private ->client [name transport init-timeout workspaces |
132 | | - {:keys [on-tools-change]}] |
| 132 | + {:keys [on-tools-change pending-tools-refresh*]}] |
133 | 133 | (let [tools-consumer (fn [tools] |
134 | 134 | (logger/info logger-tag |
135 | 135 | (format "[%s] Tools list changed, received %d tools" |
|
142 | 142 | {:name (:name %)}) |
143 | 143 | workspaces)} |
144 | 144 | :notification-handlers |
145 | | - {psd/method-notifications-tools-list_changed |
146 | | - (non-blocking-handler |
147 | | - (fn [notification] |
148 | | - (pcs/fetch-tools notification |
149 | | - {:on-tools tools-consumer}))) |
| 145 | + {;; Uses custom wrapping instead of non-blocking-handler to coordinate |
| 146 | + ;; a promise that await-pending-tools-refresh can block on. |
| 147 | + psd/method-notifications-tools-list_changed |
| 148 | + (pcs/wrap-initialized-check |
| 149 | + (fn [jsonrpc-notification] |
| 150 | + (let [p (promise)] |
| 151 | + (when pending-tools-refresh* |
| 152 | + (reset! pending-tools-refresh* p)) |
| 153 | + (future |
| 154 | + (try |
| 155 | + (pcs/fetch-tools jsonrpc-notification |
| 156 | + {:on-tools tools-consumer}) |
| 157 | + (catch Exception e |
| 158 | + (logger/error logger-tag |
| 159 | + (format "[%s] Failed to refresh tools after list_changed: %s" |
| 160 | + name (.getMessage ^Throwable e)))) |
| 161 | + (finally |
| 162 | + (when pending-tools-refresh* |
| 163 | + (compare-and-set! pending-tools-refresh* p nil)) |
| 164 | + (deliver p true))))))) |
150 | 165 |
|
151 | 166 | psd/method-notifications-resources-list_changed |
152 | 167 | (non-blocking-handler |
|
353 | 368 | server-config |
354 | 369 | {:on-server-updated on-server-updated}) |
355 | 370 | (let [init-timeout (:mcpTimeoutSeconds config) |
| 371 | + pending-tools-refresh* (atom nil) |
356 | 372 | on-tools-change (fn [tools] |
357 | 373 | (let [tools (mapv tool->internal tools)] |
358 | 374 | (swap! db* assoc-in [:mcp-clients name :tools] tools) |
|
361 | 377 | (let [{:keys [transport http-client needs-reinit?*]} (->transport name server-config workspaces db*) |
362 | 378 | result (try |
363 | 379 | (let [client (->client name transport init-timeout workspaces |
364 | | - {:on-tools-change on-tools-change}) |
| 380 | + {:on-tools-change on-tools-change |
| 381 | + :pending-tools-refresh* pending-tools-refresh*}) |
365 | 382 | init-result (pmc/get-initialize-result client) |
366 | 383 | version (get-in init-result [:serverInfo :version])] |
367 | 384 | (swap! db* assoc-in [:mcp-clients name] (cond-> {:client client |
368 | 385 | :status :starting |
369 | | - :needs-reinit?* needs-reinit?*} |
| 386 | + :needs-reinit?* needs-reinit?* |
| 387 | + :pending-tools-refresh* pending-tools-refresh*} |
370 | 388 | http-client (assoc :http-client http-client))) |
371 | 389 | (swap! db* assoc-in [:mcp-clients name :version] version) |
372 | 390 | (swap! db* assoc-in [:mcp-clients name :instructions] (:instructions init-result)) |
|
628 | 646 | :version version}) tools))) |
629 | 647 | (:mcp-clients db))) |
630 | 648 |
|
| 649 | +(defn await-pending-tools-refresh |
| 650 | + "Waits for any pending tool list refreshes to complete, with a timeout. |
| 651 | + Call before reading tools from db* to ensure dynamically loaded tools |
| 652 | + are available after a tools/list_changed notification. |
| 653 | + The pending-tools-refresh* atoms are set by the tools/list_changed |
| 654 | + notification handler and cleared when the refresh completes. |
| 655 | + Uses a shared deadline so multiple pending servers don't multiply the wait." |
| 656 | + [db timeout-ms] |
| 657 | + (let [deadline (+ (System/currentTimeMillis) timeout-ms)] |
| 658 | + (doseq [[_ {:keys [pending-tools-refresh*]}] (:mcp-clients db)] |
| 659 | + (when-let [p (and pending-tools-refresh* @pending-tools-refresh*)] |
| 660 | + (let [remaining (- deadline (System/currentTimeMillis))] |
| 661 | + (when (pos? remaining) |
| 662 | + (deref p remaining nil))))))) |
| 663 | + |
631 | 664 | (defn ^:private reinitialize-server! |
632 | 665 | "Re-initialize an MCP server after a transport error (HTTP 401/403/5xx). |
633 | 666 | Stops the old transport without attempting disconnect (the session is already |
|
0 commit comments