Fix client parsing (#8)

* Allow decoding Decodable type directly from Value

Summary: I noticed we are not converting value associated with response to the strong types requests recorded
This change will help to decode the value type later once we support the rest of plumbing (next commit)

* Store and pass original request when response arrives

Summary: Looks like there were multiple issues with parsing response:
1. It appears that we would mistakenly pass response both for response and request arguments in the private func handleResponse(_ response: Response<AnyMethod>, for request: Any)
2. Later I couldn't figure out how to win against swift type system, so ended up creating a type-erased request to store in the pendingRequests instead of original type under Any (that wouldn't allow to downcast original type later)

Test Plan:

1. Standup iMCP client from https://github.com/loopwork-ai/iMCP
2. Create a dummy project setting up MCP client against "/private/var/folders/26/8gncz37x7slbfrr95d9jcr400000gn/T/AppTranslocation/62ADEC4D-3545-4E98-A612-0E6DF52CE525/d/iMCP.app/Contents/MacOS/imcp-server" and call initialize

Note response is being parsed correctly and client.initialize() does not hang indefinitely

* Fix warning: 'Cast from Client.AnyPendingRequest to unrelated type Client.PendingRequest<Any> always fails'

* Fix warning: 'Conditional cast from NWError to NWError always succeeds'

* Uncomment code in RoundtripTests

* Formatting

* Reorganize declarations

* Inline JSONDecoder extension

* Expand roundtrip tests

---------

Co-authored-by: Sash Zats sash@zats.io <>
Co-authored-by: Mattt Zmuda <mattt@loopwork.com>
This commit is contained in:
Sash Zats 2025-03-13 14:39:46 -04:00 • committed by GitHub
commit c63a872a69
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -87,12 +87,48 @@ public actor Client {
/// The task for the message handling loop
private var task: Task<Void, Never>?
/// An error indicating a type mismatch when decoding a pending request
private struct TypeMismatchError: Swift.Error {}
/// A pending request with a continuation for the result
private struct PendingRequest<T> {
let continuation: CheckedContinuation<T, Swift.Error>
}
/// A type-erased pending request
private struct AnyPendingRequest {
private let _resume: (Result<Any, Swift.Error>) -> Void
init<T: Sendable & Decodable>(_ request: PendingRequest<T>) {
_resume = { result in
switch result {
case .success(let value):
if let typedValue = value as? T {
request.continuation.resume(returning: typedValue)
} else if let value = value as? Value,
let data = try? JSONEncoder().encode(value),
let decoded = try? JSONDecoder().decode(T.self, from: data)
{
request.continuation.resume(returning: decoded)
} else {
request.continuation.resume(throwing: TypeMismatchError())
}
case .failure(let error):
request.continuation.resume(throwing: error)
}
}
}
func resume(returning value: Any) {
_resume(.success(value))
}
func resume(throwing error: Swift.Error) {
_resume(.failure(error))
}
}
/// A dictionary of type-erased pending requests, keyed by request ID
private var pendingRequests: [ID: Any] = [:]
private var pendingRequests: [ID: AnyPendingRequest] = [:]
public init(
name: String,
@ -129,8 +165,10 @@ public actor Client {
// Attempt to decode string data as AnyResponse or AnyMessage
let decoder = JSONDecoder()
if let response = try? decoder.decode(AnyResponse.self, from: data) {
await handleResponse(response, for: response)
if let response = try? decoder.decode(AnyResponse.self, from: data),
let request = pendingRequests[response.id]
{
await handleResponse(response, for: request)
} else if let message = try? decoder.decode(AnyMessage.self, from: data) {
await handleMessage(message)
} else {
@ -158,11 +196,7 @@ public actor Client {
public func disconnect() async {
// Cancel all pending requests
for (id, request) in pendingRequests {
// We know this cast is safe because we only store PendingRequest values
if let typedRequest = request as? PendingRequest<Any> {
typedRequest.continuation.resume(
throwing: Error.internalError("Client disconnected"))
}
request.resume(throwing: Error.internalError("Client disconnected"))
pendingRequests.removeValue(forKey: id)
}
@ -220,12 +254,12 @@ public actor Client {
}
}
private func addPendingRequest<T>(
private func addPendingRequest<T: Sendable & Decodable>(
id: ID,
continuation: CheckedContinuation<T, Swift.Error>,
type: T.Type
) {
pendingRequests[id] = PendingRequest(continuation: continuation)
pendingRequests[id] = AnyPendingRequest(PendingRequest(continuation: continuation))
}
private func removePendingRequest(id: ID) {
@ -320,19 +354,18 @@ public actor Client {
// MARK: -
private func handleResponse(_ response: Response<AnyMethod>, for request: Any) async {
private func handleResponse(_ response: Response<AnyMethod>, for request: AnyPendingRequest)
async
{
await logger?.debug(
"Processing response",
metadata: ["id": "\(response.id)"])
// We know this cast is safe because we only store PendingRequest values
guard let typedRequest = request as? PendingRequest<Any> else { return }
switch response.result {
case .success(let value):
typedRequest.continuation.resume(returning: value)
request.resume(returning: value)
case .failure(let error):
typedRequest.continuation.resume(throwing: error)
request.resume(throwing: error)
}
removePendingRequest(id: response.id)