envi-code/SourceCode/Code2026/cad_mcp_frame/mcp_server.cpp

#include "StdAfx.h"
#include "KTCad.h"
#include <winsock2.h>
#include <ws2tcpip.h>
#include <thread>
#include <atomic>
#include <nlohmann/json.hpp>
#include "mcp_ConfigParser.h"
#include "mcp_utils.h"

#pragma comment(lib, "ws2_32.lib")
#define WM_MCP_EXECUTE_TASK (WM_USER + 2048)
extern HINSTANCE _hdllInstance;

#define MCP_SERVER_PORT 8080
#define MCP_SERVER_CONFIG _T("mcp_config.json")

namespace
{
	std::thread g_serverThread;
	std::atomic<bool> g_bRunning(false);
	SOCKET g_listenSocket = INVALID_SOCKET;

	HWND g_hMsgWnd = NULL;
	const wchar_t *g_szMsgClassName = L"AutoAgentMcpMsgWndClass";

	// 声明一个全局工具集对象
	mcp_Tools g_mcpTools;

	// 用于跨线程传递数据的结构体
	struct TaskData
	{
		~TaskData()
		{
			if (clientSocket != INVALID_SOCKET)
			{
				::closesocket(clientSocket);
				clientSocket = INVALID_SOCKET;
			}
		}
		SOCKET clientSocket = INVALID_SOCKET;
		std::string rawJson;
	};

	// 处理工具查询请求
	std::string HandleToolsList(int msgId)
	{
		nlohmann::json resp;
		resp["jsonrpc"] = "2.0";
		resp["id"] = msgId;

		// 直接向全局工具集要数据
		resp["result"]["tools"] = g_mcpTools.ToLLM();

		return resp.dump();
	}

	// 处理工具调用请求
	std::string HandleToolsCall(int msgId, const nlohmann::json &params)
	{
		nlohmann::json resp;
		resp["jsonrpc"] = "2.0";
		resp["id"] = msgId;

		std::string toolName = params.value("name", "");
		nlohmann::json args = params.value("arguments", nlohmann::json::object());

		// 让工具集去执行，返回 payload 可能是 {"result": ...} 也可能是 {"error": ...}
		nlohmann::json payload = g_mcpTools.Call(toolName, args);

		// 将 payload 注入到外层信封中
		resp.update(payload);

		return resp.dump();
	}

	// 此函数在 AutoCAD 主 UI 线程运行，绝对安全.
	void ExecuteTaskInMainThread(TaskData *pTask)
	{
		if (pTask == nullptr) return;

		std::string responseStr;

		// 关闭 nlohmann 的异常抛出机制，安全解析
		nlohmann::json req = nlohmann::json::parse(pTask->rawJson, nullptr, false);

		if (req.is_discarded())
		{
			// 非法 JSON 报文
			responseStr = R"({"jsonrpc": "2.0", "error": {"code": -32700, "message": "Parse error"}, "id": null})";
		}
		else
		{
			std::string method = req.value("method", "");
			int msgId = req.value("id", 0);

			if (method == "tools/list")
			{
				responseStr = HandleToolsList(msgId);
			}
			else if (method == "tools/call")
			{
				responseStr = HandleToolsCall(msgId, req.value("params", nlohmann::json::object()));
			}
			else
			{
				// 找不到方法
				nlohmann::json err;
				err["jsonrpc"] = "2.0";
				err["id"] = msgId;
				err["error"] = { {"code", -32601}, {"message", "Method not found"} };
				responseStr = err.dump();
			}
		}

		// 通过传递过来的 Socket 将结果发送回大模型端
		if (pTask->clientSocket != INVALID_SOCKET)
		{
			send(pTask->clientSocket, responseStr.c_str(), responseStr.length(), 0);
		}
	}

	// 隐式窗口的回调函数 (属于主线程)
	LRESULT CALLBACK McpWndProc(HWND hWnd, UINT message, WPARAM wParam, LPARAM lParam)
	{
		if (message == WM_MCP_EXECUTE_TASK)
		{
			TaskData *pTask = reinterpret_cast<TaskData *>(lParam);
			if (pTask)
			{
				ExecuteTaskInMainThread(pTask);
				delete pTask; // 执行完毕，清理内存
			}
			return 0;
		}
		return DefWindowProc(hWnd, message, wParam, lParam);
	}

	// TCP 监听后台线程
	void TcpServerLoop()
	{
		WSADATA wsaData;
		WSAStartup(MAKEWORD(2, 2), &wsaData);

		g_listenSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		sockaddr_in serverAddr = { 0 };
		serverAddr.sin_family = AF_INET;
		inet_pton(AF_INET, "127.0.0.1", &serverAddr.sin_addr);
		serverAddr.sin_port = htons(MCP_SERVER_PORT);

		bind(g_listenSocket, (sockaddr *)&serverAddr, sizeof(serverAddr));
		listen(g_listenSocket, SOMAXCONN);

		while (g_bRunning)
		{
			// 使用 select 模型实现带超时的轮询，确保插件卸载时线程能正常退出
			fd_set readSet;
			FD_ZERO(&readSet);
			FD_SET(g_listenSocket, &readSet);

			timeval timeout = { 1, 0 }; // 1秒超时机制
			int ret = select(0, &readSet, nullptr, nullptr, &timeout);

			if (ret > 0 && FD_ISSET(g_listenSocket, &readSet))
			{
				SOCKET clientSocket = accept(g_listenSocket, nullptr, nullptr);
				if (clientSocket != INVALID_SOCKET)
				{
					char recvBuf[2048] = { 0 };
					int bytesReceived = recv(clientSocket, recvBuf, sizeof(recvBuf) - 1, 0);
					if (bytesReceived > 0)
					{
						// 将数据拷贝到堆内存，准备派发给主线程
						TaskData *pTask = new TaskData;
						pTask->clientSocket = clientSocket; // 记录 Socket
						pTask->rawJson = recvBuf;

						// 关键转变：通过 PostMessage 发送给主线程拥有的隐式窗口
						if (g_hMsgWnd != NULL)
						{
							PostMessage(g_hMsgWnd, WM_MCP_EXECUTE_TASK, 0, reinterpret_cast<LPARAM>(pTask));
						}
						else
						{
							delete pTask;
						}
					}
				}
				else
				{
					closesocket(clientSocket);
				}
			}
		}

		closesocket(g_listenSocket);
		WSACleanup();
	}

}

namespace mcp
{
	bool close_mcp_server()
	{
		g_bRunning = false;
		if (g_serverThread.joinable())
		{
			g_serverThread.join();
		}

		if (g_hMsgWnd != NULL)
		{
			DestroyWindow(g_hMsgWnd);
			UnregisterClass(g_szMsgClassName, _hdllInstance);
			g_hMsgWnd = NULL;
		}

		g_mcpTools.clear();

		acutPrintf(_T("\n[MCP Server] The plugin has been uninstalled, and the thread has been closed.\n"));
		return true;
	}

	bool init_mcp_server()
	{
		// 1. 在主线程注册并创建 Message-Only 窗口
		WNDCLASSEX wc = { 0 };
		wc.cbSize = sizeof(WNDCLASSEX);
		wc.lpfnWndProc = McpWndProc;
		wc.hInstance = _hdllInstance;
		wc.lpszClassName = g_szMsgClassName;
		RegisterClassEx(&wc);

		g_hMsgWnd = CreateWindowEx(
			0,
			g_szMsgClassName,
			L"AutoAgentMcpMsgWindow",
			0, 0, 0, 0, 0,
			HWND_MESSAGE, // 纯消息窗口，无界面、无感知
			NULL,
			_hdllInstance,
			NULL
		);
		g_bRunning = true;
		try
		{
			g_serverThread = std::thread(TcpServerLoop);
			acutPrintf(_T("\n[MCP Server] Background TCP listening thread started (Port: 8080)...\n"));
		}
		catch (const std::system_error &e)
		{
			// 线程创建失败，重置状态标志并输出错误
			g_bRunning = false;
			// e.what() 返回的是 const char*，所以在 acutPrintf 中要用 %hs 而不是 %s 或 %ls
			acutPrintf(_T("\n[MCP Server] Fatal error: Unable to create background listening thread! Reason: %hs\n"), e.what());
		}

		// 如果监听跑起来了,则注册tools
		if (g_bRunning)
		{
			if (!mcp_ConfigParser::LoadAndRegister(mcp_utils::getModelDirPath() / MCP_SERVER_CONFIG, g_mcpTools))
			{
				acutPrintf(_T("\n[MCP Server] Warning: Failed to load tool configuration!\n[MCP Server] Automatic Shutdown.\n"));
				g_bRunning = false;
			}

			g_mcpTools.print();
		}

		return g_bRunning;
	}
}