brave · Brandon-T · Mar 16, 2024 · Mar 15, 2024 · Mar 16, 2024
diff --git a/components/ai_chat/core/browser/conversation_driver.cc b/components/ai_chat/core/browser/conversation_driver.cc
@@ -196,6 +196,38 @@ void ConversationDriver::ChangeModel(const std::string& model_key) {
   InitEngine();
 }
 
+std::string ConversationDriver::GetDefaultModel() {
+  const std::string* current_default =
+      pref_service_->GetUserPrefValue(prefs::kDefaultModelKey)->GetIfString();
+  if (current_default) {
+    return *current_default;
+  }
+
+  if (last_premium_status_ == mojom::PremiumStatus::Active ||
+      last_premium_status_ == mojom::PremiumStatus::ActiveDisconnected) {
+    return features::kAIModelsPremiumDefaultKey.Get();
+  }
+
+  current_default = pref_service_->GetDefaultPrefValue(prefs::kDefaultModelKey)
+                        ->GetIfString();
+
+  return current_default ? *current_default
+                         : features::kAIModelsDefaultKey.Get();
+}
+
+void ConversationDriver::SetDefaultModel(const std::string& model_key) {
+  DCHECK(!model_key.empty());
+  // Check that the key exists
+  auto* new_model = GetModel(model_key);
+  if (!new_model) {
+    NOTREACHED() << "No matching model found for key: " << model_key;
+    return;
+  }
+
+  pref_service_->SetDefaultPrefValue(prefs::kDefaultModelKey,
+                                     base::Value(model_key));
+}
+
 const mojom::Model& ConversationDriver::GetCurrentModel() {
   auto* model = GetModel(model_key_);
   DCHECK(model);

diff --git a/components/ai_chat/core/browser/conversation_driver.h b/components/ai_chat/core/browser/conversation_driver.h
@@ -69,6 +69,8 @@ class ConversationDriver {
   ConversationDriver& operator=(const ConversationDriver&) = delete;
 
   void ChangeModel(const std::string& model_key);
+  std::string GetDefaultModel();
+  void SetDefaultModel(const std::string& model_key);
   const mojom::Model& GetCurrentModel();
   std::vector<mojom::ModelPtr> GetModels();
   const std::vector<mojom::ConversationTurn>& GetConversationHistory();

@@ -253,7 +253,7 @@ public struct AIChatView: View {
       AIChatPaywallView(
         premiumUpgrageSuccessful: { _ in
           Task { @MainActor in
-            await model.refreshPremiumStatus()
+            await model.refreshPremiumStatusOrderCredentials()
           }
         })
     }
@@ -411,6 +411,15 @@ public struct AIChatView: View {
           isPremiumPaywallPresented = true
         },
         dismissAction: {
+          Task { @MainActor in
+            // This is needed to try to mitigate a bug in SkusSDK
+            // See: https://github.com/brave/brave-browser/issues/36824
+            // Also see the comment on the function
+            if model.premiumStatus == .active || model.premiumStatus == .activeDisconnected {
+              await model.refreshPremiumStatusOrderCredentials()
+            }
+          }
+
           if let basicModel = model.models.first(where: { $0.access == .basic }) {
             model.changeModel(modelKey: basicModel.key)
             model.retryLastRequest()

@@ -98,6 +98,7 @@ struct AIChatFeedbackToastView: View {
     )
     .containerShape(RoundedRectangle(cornerRadius: 8.0, style: .continuous))
     .padding(.horizontal)
+    .colorScheme(.light)
   }
 
   private var title: String {
@@ -127,6 +128,7 @@ struct AIChatFeedbackToastView: View {
         } label: {
           Label {
             Text(Strings.close)
+
           } icon: {
             Image(systemName: "xmark")
               .foregroundStyle(Color(braveSystemName: .primary30))
@@ -140,7 +142,7 @@ struct AIChatFeedbackToastView: View {
         } label: {
           Text(Strings.AIChat.addFeedbackActionTitle)
             .font(.subheadline)
-            .foregroundStyle(.white)
+            .foregroundStyle(Color(braveSystemName: .gray10))
         }
       }
     case .error, .submitted:

@@ -26,7 +26,7 @@ struct AIChatPageContextView: View {
           Text(Strings.AIChat.infoAboutPageContext)
             .font(.footnote)
             .foregroundStyle(Color(braveSystemName: .textTertiary))
-            .fixedSize(horizontal: true, vertical: false)
+            .fixedSize(horizontal: false, vertical: true)
 
           Button(
             action: {

@@ -183,7 +183,8 @@ public struct AIChatAdvancedSettingsView: View {
         } label: {
           LabelView(
             title: Strings.AIChat.advancedSettingsDefaultModelTitle,
-            subtitle: model.currentModel.displayName
+            subtitle: model.models.first(where: { $0.key == model.defaultAIModelKey })?.displayName
+              ?? model.currentModel.displayName
           )
         }.listRowBackground(Color(.secondaryBraveGroupedBackground))
       } header: {

@@ -5,6 +5,7 @@
 
 import BraveCore
 import BraveUI
+import Preferences
 import SwiftUI
 
 struct AIChatDefaultModelView: View {
@@ -36,7 +37,7 @@ struct AIChatDefaultModelView: View {
               if model.access == .premium, aiModel.shouldShowPremiumPrompt {
                 isPresentingPaywallPremium = true
               } else {
-                aiModel.changeModel(modelKey: model.key)
+                aiModel.defaultAIModelKey = model.key
                 dismiss()
               }
             },
@@ -53,7 +54,7 @@ struct AIChatDefaultModelView: View {
                 .frame(maxWidth: .infinity, alignment: .leading)
 
                 // If the model is selected show check
-                if model.key == aiModel.currentModel.key {
+                if model.key == aiModel.defaultAIModelKey {
                   Image(braveSystemName: "leo.check.normal")
                     .foregroundStyle(Color(braveSystemName: .textInteractive))
                     .padding(.horizontal, 4.0)
@@ -92,7 +93,7 @@ struct AIChatDefaultModelView: View {
       AIChatPaywallView(
         premiumUpgrageSuccessful: { _ in
           Task { @MainActor in
-            await aiModel.refreshPremiumStatus()
+            await aiModel.refreshPremiumStatusOrderCredentials()
           }
         })
     }

@@ -61,7 +61,7 @@ public class AIChatSubscriptionDetailModelView: ObservableObject {
     }
 
     do {
-      let credentialSummary = try await BraveSkusSDK().credentialsSummary(for: .leo)
+      let credentialSummary = try await BraveSkusSDK.shared.credentialsSummary(for: .leo)
       self.credentialSummary = credentialSummary
     } catch {
       Logger.module.error("Error Fetching Skus Credential Summary: \(error)")

@@ -84,6 +84,17 @@ public class AIChatViewModel: NSObject, ObservableObject {
     }
   }
 
+  public var defaultAIModelKey: String {
+    get {
+      return api.defaultModelKey
+    }
+
+    set {
+      objectWillChange.send()
+      api.defaultModelKey = newValue
+    }
+  }
+
   public init(
     braveCore: BraveCoreMain,
     webView: WKWebView?,
@@ -163,6 +174,10 @@ public class AIChatViewModel: NSObject, ObservableObject {
   // This function should not exist
   // We should not be refreshing credentials in this model at all!
   // This should be done in Brave-Skus-Manager once VPN moves to Skus v2
+  // Note: There is another bug where this function is required due to SkusSDK being broken
+  // See: https://github.com/brave/brave-browser/issues/36851
+  // This happens when the SkusSDK returns `activeDisconnected`, `remaining_credential_count = 0`, and `expires_at = null`
+  // Automatic credential refreshing should take care of it, but it doesn't seem to work at all
   @MainActor
   func refreshPremiumStatusOrderCredentials() async {
     await refreshPremiumStatus()
@@ -171,7 +186,7 @@ public class AIChatViewModel: NSObject, ObservableObject {
     if premiumStatus == .activeDisconnected,
       let orderId = Preferences.AIChat.subscriptionOrderId.value
     {
-      try? await BraveSkusSDK().fetchCredentials(orderId: orderId, for: .leo)
+      try? await BraveSkusSDK.shared.fetchCredentials(orderId: orderId, for: .leo)
 
       // Premium status changed after refresh
       await refreshPremiumStatus()

@@ -8,26 +8,37 @@ import Preferences
 
 extension Preferences {
   public enum AIChat {
+    /// A boolean indicating whether the user has seen the AI-Chat intro screen at least once
     public static let hasSeenIntro = Option<Bool>(
       key: "aichat.intro.hasBeenSeen",
       default: false
     )
+
+    /// The date the user's current AI-Chat subscription expires
     public static let subscriptionExpirationDate = Option<Date?>(
       key: "aichat.expiration-date",
       default: nil
     )
+
+    /// The Order-ID of the user's current AI-Chat subscription
     public static let subscriptionOrderId = Option<String?>(
       key: "aichat.order-id",
       default: nil
     )
+
+    /// A boolean indicating whether or not the user has any credentials stored locally on device
     public static let subscriptionHasCredentials = Option<Bool>(
       key: "aichat.credentials",
       default: false
     )
+
+    /// A boolean indicating whether or not the user has URL-Bar/Search-Bar auto-complete for AI-Chat
     public static let autocompleteSuggestionsEnabled = Option<Bool>(
       key: "aichat.autocompletesuggestions-enabled",
       default: true
     )
+
+    /// A boolean indicating whether or not the user has dismissed the Premium Prompt on the Feedback Form
     public static let showPremiumFeedbackAd = Option<Bool>(
       key: "aichat.show-premium-feedback-ad",
       default: true

@@ -77,7 +77,10 @@ public struct SkusOrder: Codable {
 /// A class for handling Brave Skus via SkusService
 public class BraveSkusSDK {
 
-  public init() {
+  /// Singleton instance since there can only ever be one instance of SkusService
+  public static let shared = BraveSkusSDK()
+
+  private init() {
     self.skusService = Skus.SkusServiceFactory.get(privateMode: false)
   }
 

@@ -349,7 +349,7 @@ public class BraveStoreSDK: AppStoreSDK {
     }
 
     // Create a Skus-SDK for the specified product
-    let skusSDK = BraveSkusSDK()
+    let skusSDK = BraveSkusSDK.shared
 
     // Create an order for the AppStore receipt
     // If an order already exists, refreshes the order information
@@ -378,7 +378,7 @@ public class BraveStoreSDK: AppStoreSDK {
     }
 
     // Create a Skus-SDK for the specified product
-    let skusSDK = BraveSkusSDK()
+    let skusSDK = BraveSkusSDK.shared
 
     // Create an order for the AppStore receipt
     // If an order already exists, refreshes the order information

@@ -35,6 +35,8 @@ OBJC_EXPORT
 
 @property(nonatomic, readonly) bool canShowPremiumPrompt;
 
+@property(nonatomic) NSString* defaultModelKey;
+
 @property(nonatomic) bool shouldSendPageContents;
 
 - (void)changeModel:(NSString*)modelKey;

@@ -150,6 +150,14 @@ - (void)setShouldSendPageContents:(bool)should_send {
   driver_->SetShouldSendPageContents(should_send);
 }
 
+- (NSString*)defaultModelKey {
+  return base::SysUTF8ToNSString(driver_->GetDefaultModel());
+}
+
+- (void)setDefaultModelKey:(NSString*)modelKey {
+  driver_->SetDefaultModel(base::SysNSStringToUTF8(modelKey));
+}
+
 - (void)clearConversationHistory {
   driver_->ClearConversationHistory();
 }