From 20d838491a0e32b5e6f139d1938e8e69c02a6b6c Mon Sep 17 00:00:00 2001
From: Marat Alekperov <m.alekperov@gmail.com>
Date: Fri, 9 Jan 2026 10:55:36 +0000
Subject: [PATCH] Fix original request propagation bug. Add test for checking
 request id when connecting via fallback hosts. `appendQueryItem:` adds
 request_id each time, so to avoid multiple additions I've replaced it with
 `setQueryItemNamed:`.

---
 Source/ARTNSMutableURLRequest+ARTUtils.m      | 21 +++++++--
 Source/ARTRest.m                              |  4 +-
 .../Ably/ARTNSMutableURLRequest+ARTUtils.h    |  2 +-
 .../Test Utilities/TestUtilities.swift        | 30 ++++++++++++-
 Test/AblyTests/Tests/RestClientTests.swift    | 45 ++++++++-----------
 5 files changed, 67 insertions(+), 35 deletions(-)

diff --git a/Source/ARTNSMutableURLRequest+ARTUtils.m b/Source/ARTNSMutableURLRequest+ARTUtils.m
index cfe381ed6..7a0770511 100644
--- a/Source/ARTNSMutableURLRequest+ARTUtils.m
+++ b/Source/ARTNSMutableURLRequest+ARTUtils.m
@@ -2,14 +2,27 @@
 
 @implementation NSMutableURLRequest (ARTUtils)
 
-- (void)appendQueryItem:(NSURLQueryItem *)item {
+- (void)setQueryItemNamed:(NSString *)name withValue:(NSString *)value {
     NSURLComponents *components = [NSURLComponents componentsWithURL:self.URL resolvingAgainstBaseURL:YES];
-    if(components == nil) {
+    if (components == nil) {
         return;
     }
     
-    NSMutableArray<NSURLQueryItem *> *mutableQueryItems = [NSMutableArray arrayWithArray:components.queryItems];
-    [mutableQueryItems addObject:item];
+    NSMutableArray<NSURLQueryItem *> *mutableQueryItems = [NSMutableArray arrayWithArray:components.queryItems ?: @[]];
+    
+    // Remove existing query item with the same name if it exists
+    NSMutableIndexSet *indicesToRemove = [NSMutableIndexSet indexSet];
+    for (NSUInteger i = 0; i < mutableQueryItems.count; i++) {
+        if ([mutableQueryItems[i].name isEqualToString:name]) {
+            [indicesToRemove addIndex:i];
+        }
+    }
+    [mutableQueryItems removeObjectsAtIndexes:indicesToRemove];
+    
+    // Add the new query item
+    NSURLQueryItem *newItem = [NSURLQueryItem queryItemWithName:name value:value];
+    [mutableQueryItems addObject:newItem];
+    
     components.queryItems = mutableQueryItems;
     
     NSURL *modifiedURL = components.URL;
diff --git a/Source/ARTRest.m b/Source/ARTRest.m
index b17792a3e..2e2e14325 100644
--- a/Source/ARTRest.m
+++ b/Source/ARTRest.m
@@ -368,7 +368,7 @@ - (NSString *)agentIdentifierWithWrapperSDKAgents:(nullable NSDictionary<NSStrin
                 requestId = [[randomId dataUsingEncoding:NSUTF8StringEncoding] base64EncodedStringWithOptions:0];
             }
             
-            [mutableRequest appendQueryItem:[NSURLQueryItem queryItemWithName:@"request_id" value:requestId]];
+            [mutableRequest setQueryItemNamed:@"request_id" withValue:requestId];
         }
         
         // RSC15f - reset the successed fallback host on fallbackRetryTimeout expiration
@@ -464,7 +464,7 @@ - (NSString *)agentIdentifierWithWrapperSDKAgents:(nullable NSDictionary<NSStrin
                     task = [self executeRequest:newRequest
                                       fallbacks:blockFallbacks
                                         retries:retries + 1
-                              originalRequestId:originalRequestId
+                              originalRequestId:requestId
                                wrapperSDKAgents:wrapperSDKAgents
                                      completion:callback];
                     return;
diff --git a/Source/PrivateHeaders/Ably/ARTNSMutableURLRequest+ARTUtils.h b/Source/PrivateHeaders/Ably/ARTNSMutableURLRequest+ARTUtils.h
index 31c969891..29885b588 100644
--- a/Source/PrivateHeaders/Ably/ARTNSMutableURLRequest+ARTUtils.h
+++ b/Source/PrivateHeaders/Ably/ARTNSMutableURLRequest+ARTUtils.h
@@ -9,7 +9,7 @@ NS_ASSUME_NONNULL_BEGIN
  If for some reasons new URL constructed by URLComponents is `nil`, old URL is a valiid URL for this request.
  */
 - (void)replaceHostWith:(NSString *)host;
-- (void)appendQueryItem:(NSURLQueryItem *)item;
+- (void)setQueryItemNamed:(NSString *)name withValue:(NSString *)value;
 
 @end
 
diff --git a/Test/AblyTests/Test Utilities/TestUtilities.swift b/Test/AblyTests/Test Utilities/TestUtilities.swift
index c91677ad5..e9db1a4fb 100644
--- a/Test/AblyTests/Test Utilities/TestUtilities.swift	
+++ b/Test/AblyTests/Test Utilities/TestUtilities.swift	
@@ -1541,7 +1541,35 @@ extension String {
     func replace(_ value: String, withString string: String) -> String {
         return self.replacingOccurrences(of: value, with: string, options: NSString.CompareOptions.literal, range: nil)
     }
-
+    
+    func substring(after: String) -> String? {
+        guard let range = self.range(of: after) else {
+            return nil
+        }
+        let startIndex = range.upperBound
+        return String(self[startIndex...])
+    }
+    
+    func substring(before: String) -> String? {
+        guard let range = self.range(of: before) else {
+            return nil
+        }
+        let endIndex = range.lowerBound
+        return String(self[..<endIndex])
+    }
+    
+    func substring(between: String, andString: String) -> String? {
+        guard let startRange = self.range(of: between) else {
+            return nil
+        }
+        let searchStartIndex = startRange.upperBound
+        let remainingString = String(self[searchStartIndex...])
+        guard let endRange = remainingString.range(of: andString) else {
+            return nil
+        }
+        let endIndex = endRange.lowerBound
+        return String(remainingString[..<endIndex])
+    }
 }
 
 extension ARTRealtime {
diff --git a/Test/AblyTests/Tests/RestClientTests.swift b/Test/AblyTests/Tests/RestClientTests.swift
index 1f7ac5983..61b91ca85 100644
--- a/Test/AblyTests/Tests/RestClientTests.swift
+++ b/Test/AblyTests/Tests/RestClientTests.swift
@@ -2078,39 +2078,30 @@ class RestClientTests: XCTestCase {
         let test = Test()
         let options = ARTClientOptions(key: "xxxx:xxxx")
         options.addRequestIds = true
+        let client = ARTRest(options: options)
+        let internalLog = InternalLog(clientOptions: options)
+        let mockHTTP = MockHTTP(logger: internalLog)
+        testHTTPExecutor = TestProxyHTTPExecutor(http: mockHTTP, logger: internalLog)
+        client.internal.httpExecutor = testHTTPExecutor
+        mockHTTP.setNetworkState(network: .hostUnreachable, resetAfter: 2)
+        let channel = client.channels.get(test.uniqueChannelName())
 
-        let restA = ARTRest(options: options)
-        let mockHttpExecutor = MockHTTPExecutor()
-        restA.internal.httpExecutor = mockHttpExecutor
-        waitUntil(timeout: testTimeout) { done in
-            restA.channels.get(test.uniqueChannelName()).publish(nil, data: "something") { error in
-                XCTAssertNil(error)
-                guard let url = mockHttpExecutor.requests.first?.url else {
-                    fail("No requests found")
-                    return
-                }
-                expect(url.query).to(contain("request_id"))
-                done()
-            }
-        }
-
-        mockHttpExecutor.reset()
-
-        options.addRequestIds = false
-        let restB = ARTRest(options: options)
-        restB.internal.httpExecutor = mockHttpExecutor
         waitUntil(timeout: testTimeout) { done in
-            restB.channels.get(test.uniqueChannelName()).publish(nil, data: "something") { error in
+            channel.publish(nil, data: "") { error in
                 XCTAssertNil(error)
-                XCTAssertEqual(mockHttpExecutor.requests.count, 1)
-                guard let url = mockHttpExecutor.requests.first?.url else {
-                    fail("No requests found")
-                    return
-                }
-                XCTAssertNil(url.query)
                 done()
             }
         }
+        
+        let requests = testHTTPExecutor.requests
+        XCTAssertEqual(requests.count, 3)
+        let capturedURLs = requests.map { $0.url!.absoluteString }
+        XCTAssertTrue(NSRegularExpression.match(capturedURLs.at(0), pattern: "//rest.ably.io"))
+        XCTAssertTrue(NSRegularExpression.match(capturedURLs.at(1), pattern: "//[a-e].ably-realtime.com"))
+        XCTAssertTrue(NSRegularExpression.match(capturedURLs.at(2), pattern: "//[a-e].ably-realtime.com"))
+        
+        let requestIds = capturedURLs.map { $0.substring(after: "request_id=") }
+        XCTAssertEqual(Set(requestIds).count, 1) // all requestIds are equal
     }
 
     func test__096__RestClient__request_IDs__should_remain_the_same_if_a_request_is_retried_to_a_fallback_host() {