WKURLSchemeHandler NSInternalInconsistencyException with Task cancellation

Question

I have a problem with WKURLSchemeHandler and Task cancellation and provided an example implementation below.

The problem is, that sometimes right after webView(_:stop:) is called (and "Stopping task ..." is printed) either try Task.checkCancellation() does not throw, or has already been called (I am not sure), so one of the urlSchemeTask.didReceive or didFinish can crash the app with an Exeption like this:

Stopping task <WKURLSchemeTaskImpl: 0x7fd445c209c0>
*** Terminating app due to uncaught exception 'NSInternalInconsistencyException', reason: 'This task has already been stopped'

Example implementation with comments:

import WebKit

class AsyncURLSchemeHandler: NSObject, WKURLSchemeHandler {
    private var pendingTasks = [ObjectIdentifier: TaskItem]()
    
    func webView(_ webView: WKWebView, stop urlSchemeTask: WKURLSchemeTask) {
        guard let task = pendingTasks.removeValue(forKey: urlSchemeTask.id) else { return }
        print("Stopping task \(urlSchemeTask)")
        task.stop()
    }

    func webView(_ webView: WKWebView, start urlSchemeTask: WKURLSchemeTask) {
        let task = Task { [weak self] in
            var request = urlSchemeTask.request
        
            // Do some mutation on the request
            
            do {
                try Task.checkCancellation()
                
                 // Conditionally get a URLSession 
                let session: URLSession 
                
                // Fire off the request
                let (data, response) = try await session.data(for: request)
                
                await Task.yield()
                try Task.checkCancellation()

                // Report back to the scheme task
                // Either of these !! may crash in this implementation
                urlSchemeTask.didReceive(response) // !!
                urlSchemeTask.didReceive(data) // !!
                urlSchemeTask.didFinish() // !!
                
            } catch is CancellationError {
                // Do not call didFailWithError, didFinish, or didReceive in this case
                print("Task for WKURLSchemeTask \(urlSchemeTask) has been cancelled")
            } catch {
                if !Task.isCancelled {
                    // !! This can crash, too
                    urlSchemeTask.didFailWithError(error)
                }
            }
            
            self?.pendingTasks.removeValue(forKey: urlSchemeTask.id)
        }
        
        pendingTasks[urlSchemeTask.id] = .init(urlSchemeTask: urlSchemeTask, task: task)
    }
}

private extension WKURLSchemeTask {
    var id: ObjectIdentifier {
        ObjectIdentifier(self)
    }
}

private struct TaskItem {
    enum Error: Swift.Error {
        case manualCancellation
    }
    
    let urlSchemeTask: WKURLSchemeTask
    let task: Task<Void, Never>
    
    /// Should be called when urlSchemeTask has been stopped by the system
    /// Calling anything on the urlSchemeTask afterwards would result in an exception
    func stop() {
        task.cancel()
    }
    
    /// Should be called when the urlSchemeTask should be stopped manually
    func cancel() {
        task.cancel()
        urlSchemeTask.didFailWithError(Error.manualCancellation)
    }
}

Can anyone help me to avoid these crashes?

This is a crosspost of: https://developer.apple.com/forums/thread/712430

Answer 1

To fix this i do the following:

• I wrap the pending Task state in an actor
• I use a Set<ObjectIdentifier> in that actor and checked contains instead of using task cancellation.
• I use an approach like this (https://stackoverflow.com/a/35003095/5771177) to catch the objc exception and ignore it

Most importantly:

• The main issue is out of scope of this question, I restructure the things around to fix the WebView lifecycle.