Move transport message-dispatch to base type

bandoti · bandoti · commit 606993d6dc5a · 2025-02-24T16:29:52.000-04:00
diff --git a/toolcall/handler.cpp b/toolcall/handler.cpp
@@ -102,24 +102,20 @@ void toolcall::mcp_impl::initialize() {
         tools_populating_.notify_one();
     };
 
-    transport_->subscribe("set_caps", set_caps);
-
-    mcp::initialize_request req(next_id_++);
-    transport_->send(req);
-
+    transport_->send(mcp::initialize_request(next_id_++), set_caps);
     tools_populating_.wait_for(lock, std::chrono::seconds(15), [&caps_received] { return caps_received; });
-    transport_->unsubscribe<mcp::initialize_response>("set_caps");
 
-    on_list_changed update_dirty = [this] (const mcp::tools_list_changed_notification &) {
+    on_list_changed update_dirty = [&update_dirty, this] (const mcp::tools_list_changed_notification &) {
         tool_list_dirty_ = true;
+        transport_->subscribe("notifications/tools/list_changed", update_dirty);
     };
 
     bool has_tools = false;
     for (const auto & cap : caps) {
         if (cap.name == "tools") {
             has_tools = true;
             if (cap.listChanged) {
-                transport_->subscribe("update_dirty", update_dirty);
+                transport_->subscribe("notifications/tools/list_changed", update_dirty);
             }
             break;
         }
@@ -167,28 +163,22 @@ std::string toolcall::mcp_impl::tool_list() {
         std::unique_lock<std::mutex> lock(tools_mutex_);
 
         mcp::tools_list tools;
-        on_response set_tools = [this, &tools] (const mcp::tools_list_response & resp) {
+        on_response set_tools = [this, &tools, &set_tools] (const mcp::tools_list_response & resp) {
             std::unique_lock<std::mutex> lock(tools_mutex_);
 
             tools.insert(tools.end(), resp.tools().begin(), resp.tools().end());
             auto cursor = resp.next_cursor();
             if (! cursor.empty()) {
-                mcp::tools_list_request req(next_id_++, cursor);
-                transport_->send(req);
+                transport_->send(mcp::tools_list_request(next_id_++, cursor), set_tools);
                 return;
             }
             tool_list_dirty_ = false;
             lock.unlock();
             tools_populating_.notify_one();
         };
 
-        transport_->subscribe("set_tools", set_tools);
-
-        mcp::tools_list_request req(next_id_++);
-        transport_->send(req);
-
+        transport_->send(mcp::tools_list_request(next_id_++), set_tools);
         tools_populating_.wait_for(lock, std::chrono::seconds(15), [this] { return ! tool_list_dirty_; });
-        transport_->unsubscribe<mcp::tools_list_response>("set_tools");
 
         tools_ = tools_list_to_oai_json(tools);
     }
diff --git a/toolcall/mcp_messages.cpp b/toolcall/mcp_messages.cpp
@@ -1,6 +1,7 @@
 #include "mcp_messages.h"
 #include <iostream>
 #include <log.h>
+#include <stdexcept>
 
 using json = nlohmann::json;
 
@@ -280,6 +281,14 @@ mcp::tools_list_response mcp::tools_list_response::fromJson(const nlohmann::json
     return tools_list_response(j["id"], std::move(tools), next_cursor);
 }
 
+mcp::tools_list_changed_notification mcp::tools_list_changed_notification::fromJson(const nlohmann::json & j) {
+    if (! (j.is_object() && j.contains("method") &&
+           j["method"] == "notifications/tools/list_changed")) {
+        throw std::invalid_argument("Invalid tools_list_changed message");
+    }
+    return tools_list_changed_notification();
+}
+
 mcp::tools_call_request::tools_call_request(nlohmann::json id, std::string name, tool_arg_list args)
     : request(id, "tools/call"), name_(std::move(name)), args_(std::move(args))
 {
@@ -355,27 +364,3 @@ void mcp::tools_call_response::refreshResult() {
     result["content"] = content;
     this->result(std::move(result));
 }
-
-static bool has_initialized_response(const nlohmann::json & data) {
-    return data["result"].contains("capabilities");
-}
-
-static bool has_tools_list_response(const nlohmann::json & data) {
-    return data["result"].contains("tools");
-}
-
-bool mcp::create_message(const std::string & data, mcp::message_variant & message) {
-    json j = json::parse(data);
-
-    if (has_initialized_response(j)) {
-        message = mcp::initialize_response::fromJson(j);
-
-    } else if (has_tools_list_response(j)) {
-        message = mcp::tools_list_response::fromJson(j);
-
-    } else {
-        message = std::monostate();
-        return false;
-    }
-    return true;
-}
diff --git a/toolcall/mcp_messages.h b/toolcall/mcp_messages.h
@@ -220,6 +220,8 @@ namespace mcp
     public:
         tools_list_changed_notification()
             : notification("notifications/tools/list_changed") {}
+
+        static tools_list_changed_notification fromJson(const nlohmann::json & j);
     };
 
     struct tool_arg {
@@ -270,16 +272,4 @@ namespace mcp
         bool error_;
     };
 
-    using message_variant =
-        std::variant<std::monostate,
-                     initialize_request,
-                     initialize_response,
-                     initialized_notification,
-                     tools_list_request,
-                     tools_list_response,
-                     tools_list_changed_notification,
-                     tools_call_request,
-                     tools_call_response>;
-
-    bool create_message(const std::string & data, message_variant & message);
 }
diff --git a/toolcall/mcp_sse_transport.cpp b/toolcall/mcp_sse_transport.cpp
@@ -162,10 +162,12 @@ void toolcall::mcp_sse_transport::on_endpoint_event() {
 }
 
 void toolcall::mcp_sse_transport::on_message_event() {
-    mcp::message_variant message;
-    if (mcp::create_message(event_.data, message)) {
-        notify_if<mcp::initialize_response>(message);
-        notify_if<mcp::tools_list_response>(message);
+    try {
+        nlohmann::json message = nlohmann::json::parse(event_.data);
+        notify(message);
+
+    } catch (const nlohmann::json::exception & err) {
+        LOG_WRN("SSE: Invalid message \"%s\" received: \"%s\"\n", event_.data.c_str(), err.what());
     }
 }
 
diff --git a/toolcall/mcp_transport.h b/toolcall/mcp_transport.h
@@ -31,39 +31,52 @@ namespace toolcall
             map.erase(key);
         }
 
-        template <typename T>
-        void notify(const T & message) const {
-            const auto& map =
-                std::get<std::map<std::string, toolcall::callback<T>>>(
-                    subscribers_);
+        void notify(const nlohmann::json & message) {
+            std::string key;
+            if (message.contains("id")) {
+                key = message["id"].dump();
 
-            for (const auto & pair : map) {
-                pair.second(message);
-            }
-        }
+            } else if (message.contains("method")) {
+                key = message["method"].dump();
 
-        template <typename T>
-        void notify_if(const mcp::message_variant & message) {
-            if (std::holds_alternative<T>(message)) {
-                notify(std::get<T>(message));
+            } else {
+                return;
             }
+            std::apply([&key, &message, this](auto&... maps) {
+                (..., [&] {
+                    auto it = maps.find(key);
+                    if (it != maps.end()) {
+                        using callback_type = decltype(it->second);
+                        using T = typename std::decay<typename callback_type::argument_type>::type;
+
+                        it->second(T::fromJson(message));
+                        maps.erase(it);
+                    }
+                }());
+            }, subscribers_);
         }
 
     private:
         std::tuple<std::map<std::string, toolcall::callback<MessageTypes>>...> subscribers_;
     };
 
-    class mcp_transport : public mcp_message_observer<mcp::initialize_request,
-                                                      mcp::initialize_response,
-                                                      mcp::initialized_notification,
-                                                      mcp::tools_list_request,
+    class mcp_transport : public mcp_message_observer<mcp::initialize_response,
                                                       mcp::tools_list_response,
                                                       mcp::tools_list_changed_notification> {
     public:
         virtual ~mcp_transport() = default;
 
-        template <typename T>
-        bool send(const T & message) {
+        template <typename Req, typename Resp>
+        bool send(const Req & message, callback<Resp> on_response) {
+            if (message.id().has_value()) {
+                std::string id = message.id().value().dump();
+                subscribe(id, on_response);
+            }
+            return send(message);
+        }
+
+        template <typename Req>
+        bool send(const Req & message) {
             nlohmann::json json = message.toJson();
             return send(json.dump(-1));
         }

Original file line number	Diff line number	Diff line change
`@@ -162,10 +162,12 @@ void toolcall::mcp_sse_transport::on_endpoint_event() {`
`162`	`162`	`}`
`163`	`163`
`164`	`164`	`void toolcall::mcp_sse_transport::on_message_event() {`
`165`		`- mcp::message_variant message;`
`166`		`- if (mcp::create_message(event_.data, message)) {`
`167`		`- notify_if<mcp::initialize_response>(message);`
`168`		`- notify_if<mcp::tools_list_response>(message);`
	`165`	`+ try {`
	`166`	`+ nlohmann::json message = nlohmann::json::parse(event_.data);`
	`167`	`+ notify(message);`
	`168`	`+`
	`169`	`+ } catch (const nlohmann::json::exception & err) {`
	`170`	`+ LOG_WRN("SSE: Invalid message \"%s\" received: \"%s\"\n", event_.data.c_str(), err.what());`
`169`	`171`	`}`
`170`	`172`	`}`
`171`	`173`