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:
parent
9ebe12559e
commit
c63a872a69
1 changed files with 122 additions and 45 deletions
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Reference in a new issue