Prefill AI Chat with search query

.github/workflows/adhoc.yml

@@ -21,7 +21,7 @@ on:
 
 jobs:
   make-adhoc:
-    runs-on: macos-14-xlarge
+    runs-on: macos-15
     name: Make ad-hoc build
 
     steps:

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,8 +1711,10 @@ class MainViewController: UIViewController {
         Pixel.fire(pixel: pixel, withAdditionalParameters: pixelParameters, includedParameters: [.atb])
     }
 
-    private func openAIChat() {
-
+    private func openAIChat(_ query: URLQueryItem? = nil) {
+        if let query = query {
+            aiChatViewController.loadQuery(query)
+        }
 
         let roundedPageSheet = RoundedPageSheetContainerViewController(
             contentViewController: aiChatViewController,
@@ -2084,7 +2086,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

@@ -26,5 +26,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

@@ -102,6 +102,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,54 @@
+//
+//  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 {
+    enum Constants {
+        static let duckDuckGoHost = "duckduckgo.com"
+        static let chatQueryName = "ia"
+        static let chatQueryValue = "chat"
+    }
+
+    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 == Constants.duckDuckGoHost else {
+            return false
+        }
+
+        guard let urlComponents = URLComponents(url: self, resolvingAgainstBaseURL: false),
+              let queryItems = urlComponents.queryItems else {
+            return false
+        }
+
+        return queryItems.contains { $0.name == Constants.chatQueryName && $0.value == Constants.chatQueryValue }
+    }
+}

LocalPackages/AIChat/Tests/URLExtensionTests.swift

@@ -0,0 +1,149 @@
+//
+//  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
+
+final class URLExtensionTests: XCTestCase {
+    private enum TestURLs {
+        static let exampleDomain = "https://example.com"
+        static let duckDuckGoDomain = "https://duckduckgo.com"
+
+        static let example = "\(exampleDomain)"
+        static let exampleWithKeyOldValue = "\(exampleDomain)?key=oldValue"
+        static let exampleWithExistingQuery = "\(exampleDomain)?existingKey=existingValue"
+        static let exampleWithMultipleQueryItems = "\(exampleDomain)?key1=value1&key2=value2"
+        static let duckDuckGoChat = "\(duckDuckGoDomain)/?ia=chat"
+        static let duckDuckGoWithMissingQuery = "\(duckDuckGoDomain)/"
+        static let duckDuckGoDifferentQuery = "\(duckDuckGoDomain)/?ia=search"
+        static let duckDuckGoAdditionalQueryItems = "\(duckDuckGoDomain)/?ia=chat&other=param"
+    }
+
+    func testAddingQueryItemToEmptyURL() {
+        let url = URL(string: TestURLs.example)!
+        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: TestURLs.exampleWithKeyOldValue)!
+        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: TestURLs.exampleWithExistingQuery)!
+        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: TestURLs.exampleWithMultipleQueryItems)!
+        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: TestURLs.example)!
+        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: "\(TestURLs.example)?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: TestURLs.duckDuckGoChat) {
+            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: TestURLs.exampleWithExistingQuery) {
+            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: TestURLs.duckDuckGoWithMissingQuery) {
+            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: TestURLs.duckDuckGoDifferentQuery) {
+            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: TestURLs.duckDuckGoAdditionalQueryItems) {
+            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
+    }
+}