fix: add POST /sse handler for MCP Streamable HTTP transport (Claude.ai new protocol)

includes/MCP/Server.php

@@ -5,15 +5,15 @@ namespace AskTamMCP\MCP;
 defined('ABSPATH') || exit;
 
 /**
- * MCP Server — Streamable HTTP transport (2025-03-26 spec).
+ * MCP Server — supports both transports:
  *
- * Claude.ai uses POST /sse for all communication (not GET SSE stream).
+ * 1. Streamable HTTP (new, Claude.ai uses this):
- * Each POST carries a JSON-RPC message and gets a direct JSON response.
+ *    POST /sse  with Accept: application/json  → returns JSON-RPC response directly
- * We also keep GET /sse for legacy clients that expect SSE streams.
+ *    POST /sse  with Accept: text/event-stream → streams SSE response
  *
- * POST /sse  — Claude.ai sends JSON-RPC, gets JSON-RPC response
+ * 2. Legacy SSE (old):
- * GET  /sse  — legacy SSE stream for other MCP clients
+ *    GET  /sse  → opens SSE stream, sends endpoint event
- * POST /message — explicit message endpoint (also supported)
+ *    POST /message → handles JSON-RPC messages
  */
 class Server {
 
@@ -24,7 +24,7 @@ class Server {
     }
 
     public function registerRoutes(): void {
-        // SSE endpoint — GET for legacy stream, POST for Claude.ai streamable HTTP
+        // Primary endpoint — accepts both GET (legacy SSE) and POST (streamable HTTP)
         register_rest_route('asktam-mcp/v1', '/sse', [
             [
                 'methods'             => 'GET',
@@ -38,7 +38,7 @@ class Server {
             ],
         ]);
 
-        // Explicit message endpoint
+        // Legacy message endpoint (used by old SSE transport)
         register_rest_route('asktam-mcp/v1', '/message', [
             'methods'             => 'POST',
             'callback'            => [$this, 'handleMessage'],
@@ -53,20 +53,53 @@ class Server {
     }
 
     // -------------------------------------------------------------------------
-    // POST /sse — Streamable HTTP transport (Claude.ai 2025+)
+    // Streamable HTTP transport — POST /sse
-    // Each request is a JSON-RPC message, response is JSON-RPC
+    // Claude.ai sends JSON-RPC here, we respond with JSON or SSE stream
     // -------------------------------------------------------------------------
-    public function handleSsePost(\WP_REST_Request $request): \WP_REST_Response|\WP_Error {
+    public function handleSsePost(\WP_REST_Request $request): void {
         $auth = \AskTamMCP\Core\Auth::validate($request, \AskTamMCP\Core\Auth::SCOPE_READ);
         if (is_wp_error($auth)) {
-            return $auth;
+            status_header($auth->get_error_data()['status'] ?? 401);
+            header('Content-Type: application/json');
+            echo wp_json_encode(['error' => $auth->get_error_message()]);
+            exit;
         }
 
-        return $this->dispatchJsonRpc($request, $auth);
+        $body   = $request->get_json_params() ?? [];
+        $accept = $request->get_header('Accept') ?? '';
+
+        // Determine if client wants SSE stream or plain JSON
+        $wantsStream = str_contains($accept, 'text/event-stream');
+
+        // Handle batch requests (array of JSON-RPC calls)
+        if (isset($body[0])) {
+            $results = [];
+            foreach ($body as $msg) {
+                $result = $this->dispatchJsonRpc($msg, $auth);
+                if ($result !== null) {
+                    $results[] = $result;
+                }
+            }
+            if ($wantsStream) {
+                $this->streamJsonRpcResponses($results);
+            } else {
+                $this->jsonResponse($results);
+            }
+            return;
+        }
+
+        // Single JSON-RPC message
+        $result = $this->dispatchJsonRpc($body, $auth);
+
+        if ($wantsStream) {
+            $this->streamJsonRpcResponses($result ? [$result] : []);
+        } else {
+            $this->jsonResponse($result ?? new \stdClass());
+        }
     }
 
     // -------------------------------------------------------------------------
-    // GET /sse — Legacy SSE stream for non-Claude.ai clients
+    // Legacy SSE transport — GET /sse
     // -------------------------------------------------------------------------
     public function handleSseGet(\WP_REST_Request $request): void {
         $auth = \AskTamMCP\Core\Auth::validate($request, \AskTamMCP\Core\Auth::SCOPE_READ);
@@ -88,7 +121,6 @@ class Server {
         header('Access-Control-Allow-Origin: *');
 
         $messageEndpoint = rest_url('asktam-mcp/v1/message');
-
         $this->sseEvent('endpoint', $messageEndpoint);
         $this->sseEvent('message', wp_json_encode([
             'jsonrpc' => '2.0',
@@ -115,15 +147,19 @@ class Server {
     }
 
     // -------------------------------------------------------------------------
-    // POST /message — explicit endpoint (same as POST /sse)
+    // Legacy message handler — POST /message
     // -------------------------------------------------------------------------
     public function handleMessage(\WP_REST_Request $request): \WP_REST_Response|\WP_Error {
         $auth = \AskTamMCP\Core\Auth::validate($request, \AskTamMCP\Core\Auth::SCOPE_READ);
-        if (is_wp_error($auth)) {
+        if (is_wp_error($auth)) return $auth;
-            return $auth;
-        }
 
-        return $this->dispatchJsonRpc($request, $auth);
+        $body   = $request->get_json_params();
+        $result = $this->dispatchJsonRpc($body ?? [], $auth);
+
+        return new \WP_REST_Response(
+            $result ?? ['jsonrpc' => '2.0', 'id' => null, 'result' => null],
+            200
+        );
     }
 
     public function handleInfo(): \WP_REST_Response {
@@ -131,30 +167,15 @@ class Server {
             'name'      => 'Advanced MCP Plugin by AskTam',
             'version'   => ASKTAM_MCP_VERSION,
             'endpoint'  => rest_url('asktam-mcp/v1/sse'),
-            'transport'   => 'streamable-http',
+            'transport' => ['streamable-http', 'sse'],
-            'mcp_version' => '2025-03-26',
         ]);
     }
 
     // -------------------------------------------------------------------------
-    // Shared JSON-RPC dispatcher
+    // Core JSON-RPC dispatcher — shared by all transports
     // -------------------------------------------------------------------------
-    private function dispatchJsonRpc(\WP_REST_Request $request, array $auth): \WP_REST_Response {
+    private function dispatchJsonRpc(array $body, array $authRow): ?array {
-        $body = $request->get_json_params();
+        if (empty($body['method'])) return null;
-
-        if (empty($body)) {
-            // Claude.ai sometimes sends empty POST to initialize connection
-            // Return capabilities immediately
-            return $this->rpcResult(null, [
-                'protocolVersion' => '2024-11-05',
-                'capabilities'    => ['tools' => ['listChanged' => false]],
-                'serverInfo'      => ['name' => 'advanced-mcp-plugin', 'version' => ASKTAM_MCP_VERSION],
-            ]);
-        }
-
-        if (empty($body['method'])) {
-            return $this->rpcError(null, -32600, 'Invalid Request: missing method');
-        }
 
         $id     = $body['id'] ?? null;
         $method = $body['method'];
@@ -172,29 +193,59 @@ class Server {
                 return $this->rpcResult($id, ['tools' => $this->registry->getToolDefinitions()]);
 
             case 'tools/call':
-                return $this->handleToolCall($id, $params, $auth);
+                return $this->handleToolCall($id, $params, $authRow);
 
             case 'notifications/initialized':
             case 'notifications/cancelled':
-                return new \WP_REST_Response(null, 204);
+                return null; // notifications get no response
 
             default:
                 return $this->rpcError($id, -32601, "Method not found: {$method}");
         }
     }
 
-    private function handleToolCall(mixed $id, array $params, array $authRow): \WP_REST_Response {
+    // -------------------------------------------------------------------------
+    // Output helpers
+    // -------------------------------------------------------------------------
+    private function streamJsonRpcResponses(array $responses): void {
+        while (ob_get_level()) ob_end_clean();
+        header('Content-Type: text/event-stream; charset=UTF-8');
+        header('Cache-Control: no-cache, no-store');
+        header('X-Accel-Buffering: no');
+        header('Access-Control-Allow-Origin: *');
+
+        foreach ($responses as $response) {
+            $this->sseEvent('message', wp_json_encode($response));
+        }
+
+        // Send done marker
+        echo "event: done\ndata: \n\n";
+        flush();
+        exit;
+    }
+
+    private function jsonResponse(mixed $data): void {
+        while (ob_get_level()) ob_end_clean();
+        header('Content-Type: application/json; charset=UTF-8');
+        header('Access-Control-Allow-Origin: *');
+        echo wp_json_encode($data);
+        flush();
+        exit;
+    }
+
+    private function sseEvent(string $event, string $data): void {
+        echo "event: {$event}\ndata: {$data}\n\n";
+        flush();
+    }
+
+    private function handleToolCall(mixed $id, array $params, array $authRow): array {
         $toolName  = $params['name'] ?? '';
         $arguments = $params['arguments'] ?? [];
 
-        if (empty($toolName)) {
+        if (empty($toolName)) return $this->rpcError($id, -32602, 'Missing tool name');
-            return $this->rpcError($id, -32602, 'Missing tool name');
-        }
 
         $handler = $this->registry->getHandler($toolName);
-        if (!$handler) {
+        if (!$handler) return $this->rpcError($id, -32601, "Unknown tool: {$toolName}");
-            return $this->rpcError($id, -32601, "Unknown tool: {$toolName}");
-        }
 
         $tool     = $this->registry->getTool($toolName);
         $required = $tool['scope'] ?? 'read';
@@ -208,25 +259,18 @@ class Server {
         try {
             $result = $handler($arguments);
             return $this->rpcResult($id, [
-                'content' => [
+                'content' => [['type' => 'text', 'text' => is_string($result) ? $result : wp_json_encode($result)]],
-                    ['type' => 'text', 'text' => is_string($result) ? $result : wp_json_encode($result)],
-                ],
             ]);
         } catch (\Throwable $e) {
             return $this->rpcError($id, -32000, $e->getMessage());
         }
     }
 
-    private function rpcResult(mixed $id, mixed $result): \WP_REST_Response {
+    private function rpcResult(mixed $id, mixed $result): array {
-        return new \WP_REST_Response(['jsonrpc' => '2.0', 'id' => $id, 'result' => $result]);
+        return ['jsonrpc' => '2.0', 'id' => $id, 'result' => $result];
     }
 
-    private function rpcError(mixed $id, int $code, string $message): \WP_REST_Response {
+    private function rpcError(mixed $id, int $code, string $message): array {
-        return new \WP_REST_Response(['jsonrpc' => '2.0', 'id' => $id, 'error' => ['code' => $code, 'message' => $message]], 200);
+        return ['jsonrpc' => '2.0', 'id' => $id, 'error' => ['code' => $code, 'message' => $message]];
-    }
-
-    private function sseEvent(string $event, string $data): void {
-        echo "event: {$event}\ndata: {$data}\n\n";
-        flush();
     }
 }