Support server experimental and extension capabilities

.github/workflows/ci.yml

@@ -25,7 +25,7 @@ jobs:
           # Current leg: catches breakage from new Xcode and SDK releases.
           - os: macos-latest
             swift-version: Xcode default
-          - os: ubuntu-latest
+          - os: ubuntu-24.04
             swift-version: 6.1.0
 
     runs-on: ${{ matrix.os }}
@@ -35,13 +35,13 @@ jobs:
       - uses: actions/checkout@v4
 
       - name: Prepare Linux apt dependencies
-        if: matrix.os == 'ubuntu-latest'
+        if: matrix.os == 'ubuntu-24.04'
         run: |
           sudo -n apt-get update
           sudo -n apt-get install -y --no-install-recommends libcurl4-openssl-dev
 
       - name: Setup Swift on Linux
-        if: matrix.os == 'ubuntu-latest'
+        if: matrix.os == 'ubuntu-24.04'
         uses: vapor/swiftly-action@v0.2
         with:
           toolchain: ${{ matrix.swift-version }}

Sources/MCP/Server/Server.swift

@@ -108,6 +108,10 @@ public actor Server {
 
         /// Completions capabilities
         public var completions: Completions?
+        /// Experimental capabilities supported by the server
+        public var experimental: [String: Value]?
+        /// MCP extensions supported by the server
+        public var extensions: [String: Value]?
         /// Logging capabilities
         public var logging: Logging?
         /// Prompts capabilities
@@ -122,9 +126,13 @@ public actor Server {
             logging: Logging? = nil,
             prompts: Prompts? = nil,
             resources: Resources? = nil,
-            tools: Tools? = nil
+            tools: Tools? = nil,
+            experimental: [String: Value]? = nil,
+            extensions: [String: Value]? = nil
         ) {
             self.completions = completions
+            self.experimental = experimental
+            self.extensions = extensions
             self.logging = logging
             self.prompts = prompts
             self.resources = resources

Tests/MCPTests/ClientTests.swift

@@ -14,10 +14,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -124,10 +122,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -174,10 +170,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -230,10 +224,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -321,10 +313,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -401,10 +391,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -481,10 +469,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -521,10 +507,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -656,10 +640,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,
@@ -717,10 +699,8 @@ struct ClientTests {
 
         // Set up a task to handle the initialize response
         let initTask = Task {
-            try await Task.sleep(for: .milliseconds(10))
-            if let lastMessage = await transport.sentMessages.last,
-                let data = lastMessage.data(using: .utf8),
-                let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
+            let data = await transport.waitForFirstSentData()
+            if let request = try? JSONDecoder().decode(Request<Initialize>.self, from: data)
             {
                 let response = Initialize.response(
                     id: request.id,

Tests/MCPTests/Helpers/MockTransport.swift

@@ -17,6 +17,7 @@ actor MockTransport: Transport {
     var isConnected = false
 
     private(set) var sentData: [Data] = []
+    private var firstSendWaiters: [CheckedContinuation<Data, Never>] = []
     var sentMessages: [String] {
         return sentData.compactMap { data in
             guard let string = String(data: data, encoding: .utf8) else {
@@ -57,6 +58,20 @@ actor MockTransport: Transport {
             throw MCPError.transportError(POSIXError(.EIO))
         }
         sentData.append(message)
+        let waiters = firstSendWaiters
+        firstSendWaiters.removeAll()
+        for waiter in waiters {
+            waiter.resume(returning: message)
+        }
+    }
+
+    func waitForFirstSentData() async -> Data {
+        if let data = sentData.first {
+            return data
+        }
+        return await withCheckedContinuation { continuation in
+            firstSendWaiters.append(continuation)
+        }
     }
 
     public func receive() -> AsyncThrowingStream<Data, Swift.Error> {

Tests/MCPTests/ServerTests.swift

@@ -5,6 +5,32 @@ import Testing
 
 @Suite("Server Tests")
 struct ServerTests {
+    @Test("Server capabilities preserve experimental and extension objects")
+    func testServerCapabilityObjects() throws {
+        let capabilities = Server.Capabilities(
+            tools: .init(),
+            experimental: ["example/feature": .object(["enabled": .bool(true)])],
+            extensions: ["io.modelcontextprotocol/tasks": .object([:])]
+        )
+        let result = Initialize.Result(
+            protocolVersion: Version.latest,
+            capabilities: capabilities,
+            serverInfo: .init(name: "TestServer", version: "1.0")
+        )
+
+        let encoded = try JSONEncoder().encode(result)
+        let decoded = try JSONDecoder().decode(Initialize.Result.self, from: encoded)
+        #expect(decoded.capabilities == capabilities)
+
+        let json = try #require(JSONSerialization.jsonObject(with: encoded) as? [String: Any])
+        let encodedCapabilities = try #require(json["capabilities"] as? [String: Any])
+        let experimental = try #require(encodedCapabilities["experimental"] as? [String: Any])
+        let feature = try #require(experimental["example/feature"] as? [String: Any])
+        #expect(feature["enabled"] as? Bool == true)
+        let extensions = try #require(encodedCapabilities["extensions"] as? [String: Any])
+        #expect(extensions["io.modelcontextprotocol/tasks"] is [String: Any])
+    }
+
     @Test("Start and stop server")
     func testServerStartAndStop() async throws {
         let transport = MockTransport()