Remove producers when connections close

Author: slimslenderslacks

src/jsonrpc/server.clj

@@ -259,9 +259,9 @@
                                         (fn [arguments]
                                           (logger/trace
                                            (-> arguments
-                                               (merge 
-                                                 (db/parameter-values (:name params))
-                                                 (select-keys (-> @db* :servers (get server-id)) [:roots]))
+                                               (merge
+                                                (db/parameter-values (:name params))
+                                                (select-keys (-> @db* :servers (get server-id)) [:roots]))
                                                (json/generate-string)))))
                              :id "1"}])
                           (async/reduce conj [])
@@ -307,7 +307,7 @@
                (when (= f "registry.yaml")
                  (try
                    (db/add-refs (db/registry-refs registry))
-                   (doseq [producer @jsonrpc.state/producers]
+                   (doseq [producer (vals @jsonrpc.state/producers)]
                      (try
                        (producer/publish-tool-list-changed producer {})
                        (producer/publish-prompt-list-changed producer {})
@@ -317,7 +317,7 @@
                (when (string/ends-with? f ".md")
                  (try
                    (db/update-prompt opts (string/replace f #"\.md" "") (slurp (str "/prompts/" f)))
-                   (doseq [producer @jsonrpc.state/producers]
+                   (doseq [producer (vals @jsonrpc.state/producers)]
                      (try
                        (producer/publish-tool-list-changed producer {})
                        (producer/publish-prompt-list-changed producer {})
@@ -375,9 +375,9 @@
        :trace-level trace-level
        :keyword-function keyword
        :server-context-factory
-       (fn [server]
+       (fn [server state-id]
          (let [producer (producer/->McpProducer server db*)]
-           (swap! jsonrpc.state/producers conj producer)
+           (swap! jsonrpc.state/producers assoc state-id producer)
            {:db* db*
             :logger timbre-logger
             :producer producer
@@ -399,9 +399,8 @@
                    :out System/out}))]
      (logger/info "Starting server...")
      ;; only on lsp.server/start will the stdio channels start being used
-     (let [{:keys [producer] :as context} 
-           (assoc 
-             ((:server-context-factory server-opts) server)
-             :server-id (swap! jsonrpc.state/server-counter inc))]
+     (let [server-id (swap! jsonrpc.state/server-counter inc)
+           {:keys [producer] :as context}
+           ((:server-context-factory server-opts) server server-id)]
        [producer (lsp.server/start server context)]))))
 

src/jsonrpc/socket_server.clj

@@ -27,17 +27,14 @@
    ;; chan servers have on-close callbacks
    ;; connections have both input and output streams
    (let [server-id (swap! jsonrpc.state/server-counter inc)
-         context-factory (fn [s]
-                           (assoc 
-                             (component-factory s)
-                             :server-id server-id))
          on-close #(do 
                      (shutdown/on-connection-close server-id)
                      (.close connection)
+                     (swap! jsonrpc.state/producers dissoc server-id)
                      (logger/info (format "closed connection %s" server-id)))
          s (io-server/server (assoc opts
                                     :in connection
                                     :out connection
                                     :on-close on-close))]
-     (server/start s (context-factory s)))))
+     (server/start s (component-factory s server-id)))))
 

src/jsonrpc/state.clj

@@ -2,5 +2,5 @@
 
 (def server-counter (atom 0))
 
-(def producers (atom []))
+(def producers (atom {}))