Bugfix FXIOS-16715 Keep a translation from landing on the wrong document

translateCurrentPage resolved the tab's WKWebView, awaited the model
prewarm, then called into the page. Nothing between the await and the JS
call checked for cancellation or for the document having changed, and a
tab's WKWebView is a single object that outlives its documents. When the
prewarm was slow, which happens the first time a language pair is used,
the translation was applied to whatever page the user had navigated to.

The page now mints a document identifier that native captures before the
prewarm and hands back when starting; startTranslations refuses a request
addressed to a different document. A cancellation check after the prewarm
covers the case where the task is cancelled while the model downloads.

Navigating away is not a failure, so the request is dropped silently
rather than showing a translation error on the new page.

Author: razvanlitianu

firefox-ios/Client/Frontend/Translations/Service/TranslationsService.swift

@@ -68,11 +68,19 @@ final class TranslationsService: TranslationsServiceProtocol {
         }
         onLanguageIdentified?(pageLanguage, targetLanguage)
         let webView = try currentWebView(for: windowUUID)
+        // The tab's WKWebView outlives its documents, and prewarming can take seconds.
+        let expectedDocumentId = try await currentDocumentIdJS(on: webView)
         // Prewarm resources prior to calling the JS translation API.
         await modelsFetcher.prewarmResources(for: pageLanguage, to: targetLanguage)
+        try Task.checkCancellation()
         // Create a bridge to the translations engine.
         _ = translationsEngine.bridge(to: webView)
-        try await startTranslationsJS(on: webView, from: pageLanguage, to: targetLanguage)
+        try await startTranslationsJS(
+            on: webView,
+            from: pageLanguage,
+            to: targetLanguage,
+            expecting: expectedDocumentId
+        )
     }
 
     /// Checks whether initial translation output has been produced.
@@ -101,20 +109,50 @@ final class TranslationsService: TranslationsServiceProtocol {
         }
     }
 
+    private func currentDocumentIdJS(on webView: WKWebView) async throws -> String {
+        let failure = TranslationsServiceError.jsEvaluationFailed(
+            reason: "JS evaluation failed: currentDocumentIdJS"
+        )
+        let js = "return window.__firefox__.Translations.documentId()"
+        let result: Any?
+        do {
+            result = try await webView.callAsyncJavaScript(js, contentWorld: .defaultClient)
+        } catch {
+            logger.log(
+                "Could not read the document id: \(error.localizedDescription)",
+                level: .warning,
+                category: .translations
+            )
+            throw failure
+        }
+
+        guard let documentId = result as? String, !documentId.isEmpty else { throw failure }
+        return documentId
+    }
+
     /// Starts translations by calling into the JS bridge.
+    /// Throws `.documentChanged` when the page navigated since the translation was requested.
     private func startTranslationsJS(on webView: WKWebView,
                                      from: String,
-                                     to: String) async throws {
-        let jsArgs = "{from: \"\(from)\", to: \"\(to)\"}"
-        let js = "window.__firefox__.Translations.startTranslations(\(jsArgs))"
+                                     to: String,
+                                     expecting documentId: String) async throws {
+        let jsArgs = "{from: \"\(from)\", to: \"\(to)\", expectedDocumentId: \"\(documentId)\"}"
+        let js = "return window.__firefox__.Translations.startTranslations(\(jsArgs))"
 
+        let result: Any?
         do {
-            _ = try await webView.callAsyncJavaScript(js, contentWorld: .defaultClient)
+            result = try await webView.callAsyncJavaScript(js, contentWorld: .defaultClient)
         } catch {
             /// NOTE: It would be safe to pass in the js string directly here, but it would just add too much noise 
             /// since from and to could be any language code. We only care that startTranslationsJS failed.
             throw TranslationsServiceError.jsEvaluationFailed(reason: "JS evaluation failed: startTranslationsJS")
         }
+
+        // Anything other than an explicit answer means the bridge is broken, not a navigation.
+        guard let didStart = result as? Bool else {
+            throw TranslationsServiceError.jsEvaluationFailed(reason: "JS evaluation failed: startTranslationsJS")
+        }
+        guard didStart else { throw TranslationsServiceError.documentChanged }
     }
 
     /// Evaluates the JS hook to check whether initial translation output has been produced.

firefox-ios/Client/Frontend/Translations/Service/TranslationsServiceError.swift

@@ -5,6 +5,8 @@
 /// Errors thrown by `TranslationsService` when preconditions or WebView state are invalid.
 enum TranslationsServiceError: Error, Equatable {
     case missingWebView
+    /// The page navigated before the translation started, so the request is dropped.
+    case documentChanged
     case jsEvaluationFailed(reason: String)
     case pageLanguageDetectionFailed(description: String)
     case unknown(domain: String, code: Int)
@@ -25,6 +27,8 @@ enum TranslationsServiceError: Error, Equatable {
         switch self {
         case .missingWebView:
             return "missing_webview"
+        case .documentChanged:
+            return "document_changed"
         case .jsEvaluationFailed(let reason):
             /// reason is already a stable token like "JS evaluation failed: startTranslationsJS"
             return "js_evaluation_failed(\(reason))"

firefox-ios/Client/Frontend/Translations/TranslationsMiddleware.swift

@@ -654,6 +654,17 @@ final class TranslationsMiddleware: FeatureFlaggable, Notifiable {
         } catch {
             guard !Task.isCancelled else { return }
             let serviceError = TranslationsServiceError.fromUnknown(error)
+            // Not a failure, but the icon and the telemetry flow still have to be closed out.
+            if serviceError == .documentChanged {
+                logger.log(
+                    "Translation dropped because the page changed before it started.",
+                    level: .info,
+                    category: .translations
+                )
+                dispatchClearTranslationIcon(windowUUID: windowUUID, on: tab)
+                translationFlowIds[windowUUID] = nil
+                return
+            }
             translationsTelemetry.translationFailed(
                 translationFlowId: flowId(for: windowUUID),
                 errorType: serviceError.telemetryDescription

firefox-ios/Client/Frontend/UserContent/UserScripts/MainFrame/AtDocumentStart/TranslationsEntrypoint.js

@@ -72,9 +72,14 @@ const reportPageState = () => {
 
 window.addEventListener("pageshow", reportPageState);
 
+const documentId = () => innerWindowId;
+
 /// NOTE: This should be called to start the translation process for this document.
 /// This creates the TranslationsDocument instance that manages the translation lifecycle.
-const startTranslations = ({from, to}) => {
+const startTranslations = ({from, to, expectedDocumentId}) => {
+    if (expectedDocumentId !== innerWindowId) {
+        return false;
+    }
     pendingTranslation = { from, to };
     resetIsDone();
     const languagePair = {sourceLanguage: from, targetLanguage: to}
@@ -98,6 +103,7 @@ const startTranslations = ({from, to}) => {
         innerWindowId,
     };
     sendToEngine(message);
+    return true;
 };
 
 /// NOTE: This should be called when we teardown the translations for this document.
@@ -132,5 +138,5 @@ Object.defineProperty(window.__firefox__, "Translations", {
     enumerable: false,
     configurable: false,
     writable: false,
-    value: Object.freeze({ getLanguageSampleWhenReady, startTranslations, isDone, discardTranslations })
+    value: Object.freeze({ getLanguageSampleWhenReady, documentId, startTranslations, isDone, discardTranslations })
 });

firefox-ios/firefox-ios-tests/Tests/ClientTests/TranslationsTests/TranslationsMiddlewareTests.swift

@@ -656,6 +656,60 @@ final class TranslationsMiddlewareIntegrationTests: XCTestCase, StoreTestUtility
         XCTAssertEqual(mockStore.dispatchedActions.count, 0)
     }
 
+    func test_didSelectTargetLanguage_whenDocumentChanged_doesNotDispatchError() throws {
+        setTranslationsFeatureEnabled(enabled: true)
+        let mockTranslationsService = MockTranslationsService(
+            translateResult: .failure(TranslationsServiceError.documentChanged)
+        )
+        let subject = createSubject(translationsService: mockTranslationsService)
+        let action = TranslationLanguageSelectedAction(
+            windowUUID: .XCTestDefaultUUID,
+            targetLanguage: "de",
+            actionType: TranslationsActionType.didSelectTargetLanguage
+        )
+
+        let loadingExpectation = XCTestExpectation(description: "didStartTranslatingPage dispatched")
+        mockStore.dispatchCalled = { [weak mockStore] in
+            if (mockStore?.dispatchedActions.last?.actionType as? TranslationsActionType) == .didStartTranslatingPage {
+                loadingExpectation.fulfill()
+            }
+        }
+
+        subject.translationsProvider.legacyMiddleware(mockStore.state, action)
+        wait(for: [loadingExpectation], timeout: 1.0)
+
+        // The user navigated away before the translation started. The page they are on now was
+        // never translated, so it must not be told that a translation failed.
+        let errorExpectation = XCTestExpectation(description: "didReceiveErrorTranslating should not be dispatched")
+        errorExpectation.isInverted = true
+        mockStore.dispatchCalled = { [weak mockStore] in
+            if (mockStore?.dispatchedActions.last?.actionType as? TranslationsActionType) == .didReceiveErrorTranslating {
+                errorExpectation.fulfill()
+            }
+        }
+        wait(for: [errorExpectation], timeout: 1.0)
+
+        let errorActions = mockStore.dispatchedActions.filter {
+            ($0.actionType as? TranslationsActionType) == .didReceiveErrorTranslating
+        }
+        XCTAssertTrue(errorActions.isEmpty)
+
+        let toastActions = mockStore.dispatchedActions.filter {
+            ($0.actionType as? GeneralBrowserActionType) == .showToast
+        }
+        XCTAssertTrue(toastActions.isEmpty)
+
+        XCTAssertEqual(mockTranslationsTelemetry.translationFailedCalledCount, 0)
+
+        // Dropping the request must still take the icon off `.loading`, or the spinner never ends.
+        let tab = try XCTUnwrap(mockTabManager.selectedTab)
+        XCTAssertNotEqual(tab.translationConfiguration?.state, .loading)
+        let clearActions = mockStore.dispatchedActions.filter {
+            ($0.actionType as? TranslationsActionType) == .receivedTranslationLanguage
+        }
+        XCTAssertFalse(clearActions.isEmpty)
+    }
+
     func test_didSelectTargetLanguage_withTranslationError_dispatchToastAction() throws {
         setTranslationsFeatureEnabled(enabled: true)
         enum TestError: Error { case example }