Skip to content

Commit a0ae212

Browse files
Improvements and fixes (#223)
* feat: added http transport context to the task local execution * fix: added refresh_token to the grantTypes for DCR * chore: updated documentation
1 parent cb6a62f commit a0ae212

8 files changed

Lines changed: 345 additions & 14 deletions

File tree

README.md

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ of the MCP specification.
4141
- [Logging](#logging-1)
4242
- [Progress Tracking](#progress-tracking-1)
4343
- [Initialize Hook](#initialize-hook)
44+
- [HTTP Request Context in Handlers](#http-request-context-in-handlers)
4445
- [Graceful Shutdown](#graceful-shutdown)
4546
- [Transports](#transports)
4647
- [Authentication](#authentication)
@@ -1184,6 +1185,35 @@ try await server.start(transport: transport) { clientInfo, clientCapabilities in
11841185
}
11851186
```
11861187

1188+
### HTTP Request Context in Handlers
1189+
1190+
When a server is connected over `StatefulHTTPServerTransport` or `StatelessHTTPServerTransport`,
1191+
method handlers can observe the originating HTTP request (headers, body, path, method) via
1192+
`Server.currentHandlerContext` — a task-local set automatically before each handler runs:
1193+
1194+
```swift
1195+
await server.withMethodHandler(CallTool.self) { params in
1196+
let httpRequest = Server.currentHandlerContext?.httpContext
1197+
let authHeader = httpRequest?.headers["Authorization"]
1198+
// …use the header to scope the response, audit, etc.
1199+
return .init(content: [.text("ok")], isError: false)
1200+
}
1201+
```
1202+
1203+
`httpContext` is `nil` for transports that don't carry HTTP context (e.g. `StdioTransport`,
1204+
`InMemoryTransport`) and for handlers reached off the dispatch path.
1205+
1206+
Task-locals are not inherited by `Task.detached`. If you spawn a detached task from a handler,
1207+
capture the context up front:
1208+
1209+
```swift
1210+
let ctx = Server.currentHandlerContext
1211+
Task.detached { await doWork(with: ctx?.httpContext) }
1212+
```
1213+
1214+
Custom HTTP transports can opt in by conforming to `HTTPContextProviding` and returning the
1215+
`HTTPRequest` for a given JSON-RPC id while it is in flight.
1216+
11871217
### Graceful Shutdown
11881218

11891219
We recommend using

Sources/MCP/Base/Authorization/OAuthClientRegistrar.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,7 @@ struct OAuthClientRegistrar: Sendable {
7979
let responseTypes: [String]
8080
switch configuration.grantType {
8181
case .authorizationCode:
82-
grantTypes = [OAuthGrantTypeValue.authorizationCode]
82+
grantTypes = [OAuthGrantTypeValue.authorizationCode, OAuthGrantTypeValue.refreshToken]
8383
responseTypes = [OAuthParameterName.code]
8484
case .clientCredentials:
8585
grantTypes = [OAuthGrantTypeValue.clientCredentials]

Sources/MCP/Base/Transports/HTTPServer/HTTPServerTypes.swift

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,24 @@ struct SSEEvent: Sendable {
207207
}
208208
}
209209

210+
// MARK: - HTTP Context Providing
211+
212+
/// A transport that can surface the originating HTTP request for an in-flight JSON-RPC
213+
/// request, keyed by its JSON-RPC id.
214+
///
215+
/// The server consults this during dispatch so that registered method handlers can
216+
/// observe the HTTP request that triggered them (headers, auth, path, body) via
217+
/// ``Server/currentHandlerContext``, without changing the `withMethodHandler` signature.
218+
///
219+
/// Only JSON-RPC *requests* (with an id) are addressable — notifications have no id and
220+
/// are not correlated.
221+
public protocol HTTPContextProviding: Sendable {
222+
/// Returns the HTTP request associated with the given JSON-RPC id, if the
223+
/// transport still has it on hand. Implementations must return `nil` once the
224+
/// corresponding response has been delivered or the session has terminated.
225+
func httpRequestContext(for id: ID) async -> HTTPRequest?
226+
}
227+
210228
// MARK: - JSON-RPC Message Classification
211229

212230
/// Classifies a raw JSON-RPC message for routing purposes.

Sources/MCP/Base/Transports/HTTPServer/StatefulHTTPServerTransport.swift

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import Logging
2929
/// and receive `HTTPResponse` values to convert to your framework's native types.
3030
/// For SSE responses, the `.stream` case provides an `AsyncThrowingStream<Data, Error>`
3131
/// to pipe to the client.
32-
public actor StatefulHTTPServerTransport: Transport {
32+
public actor StatefulHTTPServerTransport: Transport, HTTPContextProviding {
3333
public nonisolated let logger: Logger
3434

3535
// MARK: - Dependencies
@@ -54,6 +54,11 @@ public actor StatefulHTTPServerTransport: Transport {
5454
/// Maps request ID → SSE stream continuation for active POST request streams.
5555
private var requestSSEContinuations: [String: AsyncThrowingStream<Data, Swift.Error>.Continuation] = [:]
5656

57+
/// Maps request ID → originating HTTP request, surfaced to handlers via
58+
/// ``Server/currentHTTPContext``. Cleared when the response is routed,
59+
/// the stream is closed, or the session terminates.
60+
private var httpRequestContexts: [String: HTTPRequest] = [:]
61+
5762
// MARK: - Standalone GET SSE stream
5863

5964
/// The standalone SSE stream continuation for server-initiated messages.
@@ -280,6 +285,7 @@ public actor StatefulHTTPServerTransport: Transport {
280285
// Create SSE stream for this request
281286
let (sseStream, sseContinuation) = AsyncThrowingStream<Data, Swift.Error>.makeStream()
282287
requestSSEContinuations[requestID] = sseContinuation
288+
httpRequestContexts[requestID] = request
283289

284290
// Extract protocol version for priming event decision
285291
let protocolVersion = extractProtocolVersion(from: body, request: request)
@@ -406,6 +412,7 @@ public actor StatefulHTTPServerTransport: Transport {
406412
continuation.finish()
407413
requestSSEContinuations.removeValue(forKey: requestID)
408414
}
415+
httpRequestContexts.removeValue(forKey: requestID)
409416
}
410417

411418
private func routeServerInitiatedMessage(_ data: Data) {
@@ -535,6 +542,13 @@ public actor StatefulHTTPServerTransport: Transport {
535542
guard let continuation = requestSSEContinuations[requestID] else { return }
536543
continuation.finish()
537544
requestSSEContinuations.removeValue(forKey: requestID)
545+
httpRequestContexts.removeValue(forKey: requestID)
546+
}
547+
548+
// MARK: - HTTPContextProviding
549+
550+
public func httpRequestContext(for id: ID) -> HTTPRequest? {
551+
httpRequestContexts[id.description]
538552
}
539553

540554
// MARK: - Termination
@@ -552,6 +566,7 @@ public actor StatefulHTTPServerTransport: Transport {
552566
continuation.finish()
553567
}
554568
requestSSEContinuations.removeAll()
569+
httpRequestContexts.removeAll()
555570

556571
// Close standalone GET stream
557572
standaloneSSEContinuation?.finish()

Sources/MCP/Base/Transports/HTTPServer/StatelessHTTPServerTransport.swift

Lines changed: 22 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ import Logging
2929
/// - Session management is handled externally or not needed
3030
///
3131
/// For full streaming and session support, use ``StatefulHTTPServerTransport`` instead.
32-
public actor StatelessHTTPServerTransport: Transport {
32+
public actor StatelessHTTPServerTransport: Transport, HTTPContextProviding {
3333
public nonisolated let logger: Logger
3434

3535
// MARK: - Dependencies
@@ -52,6 +52,11 @@ public actor StatelessHTTPServerTransport: Transport {
5252
/// When the server calls `send()` with a response, the matching continuation is resumed.
5353
private var responseWaiters: [String: CheckedContinuation<Data, any Error>] = [:]
5454

55+
/// Maps request ID → originating HTTP request, surfaced to handlers via
56+
/// ``Server/currentHTTPContext``. Entries live only while a JSON-RPC request
57+
/// is in flight.
58+
private var httpRequestContexts: [String: HTTPRequest] = [:]
59+
5560
// MARK: - Init
5661

5762
/// Creates a new stateless HTTP server transport.
@@ -207,11 +212,16 @@ public actor StatelessHTTPServerTransport: Transport {
207212
return .accepted()
208213

209214
case .request(let id, _):
210-
return await handleJSONRPCRequest(body, requestID: id)
215+
return await handleJSONRPCRequest(body, requestID: id, request: request)
211216
}
212217
}
213218

214-
private func handleJSONRPCRequest(_ body: Data, requestID: String) async -> HTTPResponse {
219+
private func handleJSONRPCRequest(
220+
_ body: Data,
221+
requestID: String,
222+
request: HTTPRequest
223+
) async -> HTTPResponse {
224+
httpRequestContexts[requestID] = request
215225
// Yield the incoming message to the server
216226
incomingContinuation.yield(body)
217227

@@ -222,15 +232,23 @@ public actor StatelessHTTPServerTransport: Transport {
222232
responseWaiters[requestID] = continuation
223233
}
224234
} catch {
235+
httpRequestContexts.removeValue(forKey: requestID)
225236
return .error(
226237
statusCode: 500,
227238
.internalError("Error processing request: \(error.localizedDescription)")
228239
)
229240
}
230241

242+
httpRequestContexts.removeValue(forKey: requestID)
231243
return .data(responseData, headers: [HTTPHeaderName.contentType: ContentType.json])
232244
}
233245

246+
// MARK: - HTTPContextProviding
247+
248+
public func httpRequestContext(for id: ID) -> HTTPRequest? {
249+
httpRequestContexts[id.description]
250+
}
251+
234252
// MARK: - Termination
235253

236254
private func terminate() async {
@@ -245,6 +263,7 @@ public actor StatelessHTTPServerTransport: Transport {
245263
logger.debug("Cancelled waiter for request", metadata: ["requestID": "\(id)"])
246264
}
247265
responseWaiters.removeAll()
266+
httpRequestContexts.removeAll()
248267

249268
// Close incoming stream
250269
incomingContinuation.finish()

Sources/MCP/Server/Server.swift

Lines changed: 42 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -298,13 +298,41 @@ public actor Server {
298298

299299
// MARK: - Request Context
300300

301-
/// The JSON-RPC request ID of the currently executing method handler.
301+
/// Per-dispatch context the SDK attaches to the currently executing handler.
302+
///
303+
/// Exposed via the ``Server/currentHandlerContext`` task-local so method
304+
/// handlers can observe the originating HTTP request (headers, auth, path,
305+
/// body) without changing the `withMethodHandler` signature.
306+
public struct HandlerContext: Sendable {
307+
/// The JSON-RPC request id of the in-flight request. SDK-internal use
308+
/// (e.g. transports closing an SSE stream mid-call per SEP-1699).
309+
package let id: ID
310+
311+
/// The originating HTTP request, if the active transport conforms to
312+
/// ``HTTPContextProviding``. `nil` for transports that don't carry HTTP
313+
/// context (stdio, in-memory) or for handlers reached off the dispatch
314+
/// path.
315+
public let httpContext: HTTPRequest?
316+
317+
package init(id: ID, httpContext: HTTPRequest?) {
318+
self.id = id
319+
self.httpContext = httpContext
320+
}
321+
}
322+
323+
/// The handler context for the currently executing method handler.
302324
///
303325
/// Set via `@TaskLocal` before dispatching each request, so it propagates
304-
/// automatically into the handler task. Accessible package-wide for
305-
/// transports that need to identify the active request (e.g. closing an
306-
/// SSE stream mid-call for reconnection testing per SEP-1699).
307-
@TaskLocal package static var currentRequestID: ID? = nil
326+
/// automatically into the handler task. `nil` outside of a handler.
327+
///
328+
/// When spawning `Task.detached { … }` from inside a handler, capture the
329+
/// value up front — detached tasks do not inherit task-locals:
330+
///
331+
/// ```swift
332+
/// let ctx = Server.currentHandlerContext
333+
/// Task.detached { await doWork(with: ctx?.httpContext) }
334+
/// ```
335+
@TaskLocal public static var currentHandlerContext: HandlerContext? = nil
308336

309337
// MARK: - Registration
310338

@@ -750,10 +778,17 @@ public actor Server {
750778
return response
751779
}
752780

781+
// Ask the transport for the originating HTTP request so handlers can observe
782+
// headers/auth via `Server.currentHandlerContext?.httpContext`. Transports
783+
// that don't carry HTTP context (stdio, in-memory) don't conform.
784+
let httpContext = await (connection as? any HTTPContextProviding)?
785+
.httpRequestContext(for: request.id)
786+
let handlerContext = HandlerContext(id: request.id, httpContext: httpContext)
787+
753788
// Create a task to handle the request with cancellation support.
754-
// Set currentRequestID as a task local so handlers can identify the active request.
789+
// Set currentHandlerContext as a task local so handlers see it.
755790
var handlerTask: Task<Response<AnyMethod>, Error>!
756-
Server.$currentRequestID.withValue(request.id) {
791+
Server.$currentHandlerContext.withValue(handlerContext) {
757792
handlerTask = Task<Response<AnyMethod>, Error> {
758793
do {
759794
// Check if task was cancelled before starting

Sources/MCPConformance/Server/main.swift

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -180,8 +180,7 @@ func createConformanceServer(state: ServerState, transport: StatefulHTTPServerTr
180180
// SEP-1699: Close the SSE stream mid-call to trigger client reconnection.
181181
// The client should reconnect via GET with Last-Event-ID and receive the
182182
// response on the new stream.
183-
if let requestID = Server.currentRequestID {
184-
183+
if let requestID = Server.currentHandlerContext?.id {
185184
await transport.closeSSEStream(forRequestID: requestID.description)
186185
}
187186
// Wait briefly for the client to reconnect before sending the response.

0 commit comments

Comments
 (0)