Prefill AI Chat with search query

DuckDuckGo.xcodeproj/xcshareddata/xcschemes/DuckDuckGo.xcscheme

@@ -110,6 +110,16 @@
                ReferencedContainer = "container:DuckDuckGo.xcodeproj">
             </BuildableReference>
          </TestableReference>
+         <TestableReference
+            skipped = "NO">
+            <BuildableReference
+               BuildableIdentifier = "primary"
+               BlueprintIdentifier = "AIChatTests"
+               BuildableName = "AIChatTests"
+               BlueprintName = "AIChatTests"
+               ReferencedContainer = "container:LocalPackages/AIChat">
+            </BuildableReference>
+         </TestableReference>
       </Testables>
    </TestAction>
    <LaunchAction

DuckDuckGo/MainViewController.swift

@@ -1711,10 +1711,14 @@ class MainViewController: UIViewController {
         Pixel.fire(pixel: pixel, withAdditionalParameters: pixelParameters, includedParameters: [.atb])
     }
 
-    private func openAIChat() {
+    private func openAIChat(_ query: URLQueryItem? = nil) {
         let logoImage = UIImage(named: "Logo")
         let title = UserText.aiChatTitle
 
+        if let query = query {
+            aiChatViewController.loadQuery(query)
+        }
+
         let roundedPageSheet = RoundedPageSheetContainerViewController(
             contentViewController: aiChatViewController,
             logoImage: logoImage,
@@ -2087,7 +2091,8 @@ extension MainViewController: OmniBarDelegate {
 
         switch accessoryType {
         case .chat:
-            openAIChat()
+            let queryItem = currentTab?.url?.getQueryItems()?.filter { $0.name == "q" }.first
+            openAIChat(queryItem)
             Pixel.fire(pixel: .openAIChatFromAddressBar)
         case .share:
             Pixel.fire(pixel: .addressBarShare)

LocalPackages/AIChat/Package.swift

@@ -20,5 +20,9 @@ let package = Package(
                 .process("Resources/Assets.xcassets")
             ]
         ),
+        .testTarget(
+            name: "AIChatTests",
+            dependencies: ["AIChat"]
+        )
     ]
 )

LocalPackages/AIChat/Sources/AIChat/AIChatWebViewController.swift

@@ -96,6 +96,11 @@ extension AIChatWebViewController {
         let request = URLRequest(url: chatModel.aiChatURL)
         webView.load(request)
     }
+
+    func loadQuery(_ query: URLQueryItem) {
+        let queryURL = chatModel.aiChatURL.addingOrReplacingQueryItem(query)
+        webView.load(URLRequest(url: queryURL))
+    }
 }
 
 // MARK: - WKNavigationDelegate
@@ -104,7 +109,7 @@ extension AIChatWebViewController: WKNavigationDelegate {
 
     func webView(_ webView: WKWebView, decidePolicyFor navigationAction: WKNavigationAction) async -> WKNavigationActionPolicy {
         if let url = navigationAction.request.url {
-            if url == chatModel.aiChatURL || navigationAction.targetFrame?.isMainFrame == false {
+            if url.isDuckAIURL || navigationAction.targetFrame?.isMainFrame == false {
                 return .allow
             } else {
                 delegate?.aiChatWebViewController(self, didRequestToLoad: url)

LocalPackages/AIChat/Sources/AIChat/Public API/AIChatViewController.swift

@@ -84,6 +84,17 @@ extension AIChatViewController {
     }
 }
 
+// MARK: - Public functions
+extension AIChatViewController {
+    public func loadQuery(_ query: URLQueryItem) {
+         // Ensure the webViewController is added before loading the query
+         if webViewController == nil {
+             addWebViewController()
+         }
+         webViewController?.loadQuery(query)
+     }
+}
+
 // MARK: - Views Setup
 extension AIChatViewController {
 

LocalPackages/AIChat/Sources/AIChat/URL+Extension.swift

@@ -0,0 +1,48 @@
+//
+//  URL+Extension.swift
+//  DuckDuckGo
+//
+//  Copyright © 2024 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import Foundation
+
+extension URL {
+    func addingOrReplacingQueryItem(_ queryItem: URLQueryItem) -> URL {
+        guard var urlComponents = URLComponents(url: self, resolvingAgainstBaseURL: false) else {
+            return self
+        }
+
+        var queryItems = urlComponents.queryItems ?? []
+        queryItems.removeAll { $0.name == queryItem.name }
+        queryItems.append(queryItem)
+
+        urlComponents.queryItems = queryItems
+        return urlComponents.url ?? self
+    }
+
+    var isDuckAIURL: Bool {
+        guard let host = self.host, host == "duckduckgo.com" else {
+            return false
+        }
+
+        guard let urlComponents = URLComponents(url: self, resolvingAgainstBaseURL: false),
+              let queryItems = urlComponents.queryItems else {
+            return false
+        }
+
+        return queryItems.contains { $0.name == "ia" && $0.value == "chat" }
+    }
+}

LocalPackages/AIChat/Tests/URLExtensionTests.swift

@@ -0,0 +1,136 @@
+//
+//  URLExtensionTests.swift
+//  DuckDuckGo
+//
+//  Copyright © 2022 DuckDuckGo. All rights reserved.
+//
+//  Licensed under the Apache License, Version 2.0 (the "License");
+//  you may not use this file except in compliance with the License.
+//  You may obtain a copy of the License at
+//
+//  http://www.apache.org/licenses/LICENSE-2.0
+//
+//  Unless required by applicable law or agreed to in writing, software
+//  distributed under the License is distributed on an "AS IS" BASIS,
+//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+//  See the License for the specific language governing permissions and
+//  limitations under the License.
+//
+
+import XCTest
+@testable import AIChat
+
+class URLExtensionTests: XCTestCase {
+
+    func testAddingQueryItemToEmptyURL() {
+        let url = URL(string: "https://example.com")!
+        let queryItem = URLQueryItem(name: "key", value: "value")
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["key": "value"])
+    }
+
+    func testReplacingExistingQueryItem() {
+        let url = URL(string: "https://example.com?key=oldValue")!
+        let queryItem = URLQueryItem(name: "key", value: "newValue")
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["key": "newValue"])
+    }
+
+    func testAddingQueryItemToExistingQuery() {
+        let url = URL(string: "https://example.com?existingKey=existingValue")!
+        let queryItem = URLQueryItem(name: "newKey", value: "newValue")
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["existingKey": "existingValue", "newKey": "newValue"])
+    }
+
+    func testReplacingOneOfMultipleQueryItems() {
+        let url = URL(string: "https://example.com?key1=value1&key2=value2")!
+        let queryItem = URLQueryItem(name: "key1", value: "newValue1")
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["key1": "newValue1", "key2": "value2"])
+    }
+
+    func testAddingQueryItemWithNilValue() {
+        let url = URL(string: "https://example.com")!
+        let queryItem = URLQueryItem(name: "key", value: nil)
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["key": ""])
+    }
+
+    func testReplacingQueryItemWithNilValue() {
+        let url = URL(string: "https://example.com?key=value")!
+        let queryItem = URLQueryItem(name: "key", value: nil)
+        let result = url.addingOrReplacingQueryItem(queryItem)
+
+        XCTAssertEqual(result.scheme, "https")
+        XCTAssertEqual(result.host, "example.com")
+        XCTAssertEqual(result.queryItemsDictionary, ["key": ""])
+    }
+
+    func testIsDuckAIURLWithValidURL() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=chat") {
+            XCTAssertTrue(url.isDuckAIURL, "The URL should be identified as a DuckDuckGo AI URL.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithInvalidDomain() {
+        if let url = URL(string: "https://example.com/?ia=chat") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to the domain.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithMissingQueryItem() {
+        if let url = URL(string: "https://duckduckgo.com/") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to missing query item.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithDifferentQueryItem() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=search") {
+            XCTAssertFalse(url.isDuckAIURL, "The URL should not be identified as a DuckDuckGo AI URL due to different query item value.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+
+    func testIsDuckAIURLWithAdditionalQueryItems() {
+        if let url = URL(string: "https://duckduckgo.com/?ia=chat&other=param") {
+            XCTAssertTrue(url.isDuckAIURL, "The URL should be identified as a DuckDuckGo AI URL even with additional query items.")
+        } else {
+            XCTFail("Failed to create URL from string.")
+        }
+    }
+}
+
+extension URL {
+    var queryItemsDictionary: [String: String] {
+        var dict = [String: String]()
+        if let queryItems = URLComponents(url: self, resolvingAgainstBaseURL: false)?.queryItems {
+            for item in queryItems {
+                dict[item.name] = item.value ?? ""
+            }
+        }
+        return dict
+    }
+}