From e79f1811dac7627f5446fa8cca3788f68588013c Mon Sep 17 00:00:00 2001
From: Zekun Jiang <[email protected]>
Date: Fri, 25 Sep 2026 13:15:28 -0700
Subject: [PATCH] Remove storage of raw NetworkContext pointers across private_ai.

Bug: 558464760
Change-Id: I5a28e778b73dc797aac3e2ce6de5a4ef328d683b
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/8445033
Reviewed-by: Christian Dullweber <[email protected]>
Commit-Queue: Zekun Jiang <[email protected]>
Reviewed-by: Rohit Rao <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1705101}
---
diff --git a/chrome/browser/private_ai/private_ai_service_factory.cc b/chrome/browser/private_ai/private_ai_service_factory.cc
index 16edb01..0e6c0d8 100644
--- a/chrome/browser/private_ai/private_ai_service_factory.cc
+++ b/chrome/browser/private_ai/private_ai_service_factory.cc
@@ -4,6 +4,7 @@
 
 #include "chrome/browser/private_ai/private_ai_service_factory.h"
 
+#include "base/functional/bind.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/profiles/profile_selections.h"
 #include "chrome/browser/signin/identity_manager_factory.h"
@@ -29,6 +30,18 @@
   return instance.get();
 }
 
+// static
+network::NetworkContextGetter
+PrivateAiServiceFactory::CreateNetworkContextGetter(Profile* profile) {
+  return base::BindRepeating(
+      [](base::WeakPtr<Profile> profile) -> network::mojom::NetworkContext* {
+        return profile
+                   ? profile->GetDefaultStoragePartition()->GetNetworkContext()
+                   : nullptr;
+      },
+      profile->GetWeakPtr());
+}
+
 PrivateAiServiceFactory::PrivateAiServiceFactory()
     : ProfileKeyedServiceFactory("private_ai::PrivateAiService",
                                  ProfileSelections::BuildForRegularProfile()) {
@@ -57,9 +70,8 @@
       profile->GetDefaultStoragePartition()
           ->GetURLLoaderFactoryForBrowserProcess(),
       std::move(network_driver), std::move(oak_session_driver),
-      profile->GetDefaultStoragePartition()->GetNetworkContext(),
-      kPrivateAiUrl.Get(), PrivateAiService::GetApiKey(channel),
-      kPrivateAiProxyServerUrl.Get(),
+      CreateNetworkContextGetter(profile), kPrivateAiUrl.Get(),
+      PrivateAiService::GetApiKey(channel), kPrivateAiProxyServerUrl.Get(),
       base::FeatureList::IsEnabled(kPrivateAiUseTokenAttestation), channel);
 }
 
diff --git a/chrome/browser/private_ai/private_ai_service_factory.h b/chrome/browser/private_ai/private_ai_service_factory.h
index 726da633..b35fc13 100644
--- a/chrome/browser/private_ai/private_ai_service_factory.h
+++ b/chrome/browser/private_ai/private_ai_service_factory.h
@@ -8,6 +8,7 @@
 #include "base/no_destructor.h"
 #include "chrome/browser/profiles/profile_keyed_service_factory.h"
 #include "chrome/browser/profiles/profile_selections.h"
+#include "services/network/public/cpp/network_context_getter.h"
 
 class Profile;
 
@@ -19,6 +20,8 @@
  public:
   static PrivateAiService* GetForProfile(Profile* profile);
   static PrivateAiServiceFactory* GetInstance();
+  static network::NetworkContextGetter CreateNetworkContextGetter(
+      Profile* profile);
 
   PrivateAiServiceFactory(const PrivateAiServiceFactory&) = delete;
   PrivateAiServiceFactory& operator=(const PrivateAiServiceFactory&) = delete;
diff --git a/chrome/browser/ui/webui/private_ai_internals/private_ai_internals_ui.cc b/chrome/browser/ui/webui/private_ai_internals/private_ai_internals_ui.cc
index 8b8f8fd..1b40942 100644
--- a/chrome/browser/ui/webui/private_ai_internals/private_ai_internals_ui.cc
+++ b/chrome/browser/ui/webui/private_ai_internals/private_ai_internals_ui.cc
@@ -22,7 +22,6 @@
 #include "components/version_info/channel.h"
 #include "content/public/browser/browser_context.h"
 #include "content/public/browser/internal_webui_config.h"
-#include "content/public/browser/storage_partition.h"
 #include "content/public/browser/web_ui.h"
 #include "content/public/browser/web_ui_data_source.h"
 #include "mojo/public/cpp/bindings/pending_receiver.h"
@@ -85,12 +84,12 @@
   auto* token_manager = private_ai_service->GetTokenManager();
   CHECK(token_manager);
 
-  auto* network_context =
-      profile->GetDefaultStoragePartition()->GetNetworkContext();
   page_handler_ = std::make_unique<PrivateAiInternalsPageHandler>(
-      token_manager, network_context, private_ai_service->GetClient(),
-      private_ai_service->GetLogger(), &oak_session_driver_content_,
-      &network_driver_content_, chrome::GetChannel(), std::move(receiver));
+      token_manager,
+      PrivateAiServiceFactory::CreateNetworkContextGetter(profile),
+      private_ai_service->GetClient(), private_ai_service->GetLogger(),
+      &oak_session_driver_content_, &network_driver_content_,
+      chrome::GetChannel(), std::move(receiver));
 }
 
 WEB_UI_CONTROLLER_TYPE_IMPL(PrivateAiInternalsUI)
diff --git a/components/private_ai/client.cc b/components/private_ai/client.cc
index 0653d04..77c362e1 100644
--- a/components/private_ai/client.cc
+++ b/components/private_ai/client.cc
@@ -17,7 +17,7 @@
 #include "components/private_ai/features.h"
 #include "components/private_ai/phosphor/token_manager.h"
 #include "net/base/url_util.h"
-#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "url/gurl.h"
 #include "url/url_constants.h"
 
@@ -29,7 +29,7 @@
     const std::string& api_key,
     const std::string& proxy_url_string,
     bool use_token_attestation,
-    network::mojom::NetworkContext* network_context,
+    network::NetworkContextGetter network_context_getter,
     phosphor::TokenManager* token_manager,
     PrivateAiLogger* logger,
     PrivateAiOakSessionDriver* oak_session_driver,
@@ -39,8 +39,8 @@
   GURL formatted_url = Client::FormatUrl(url, api_key);
 
   auto connection_factory = std::make_unique<ConnectionFactoryImpl>(
-      formatted_url, network_context, logger, oak_session_driver,
-      network_driver, channel);
+      formatted_url, std::move(network_context_getter), logger,
+      oak_session_driver, network_driver, channel);
 
   if (use_token_attestation) {
     connection_factory->EnableTokenAttestation(token_manager);
diff --git a/components/private_ai/client.h b/components/private_ai/client.h
index b4e57f876..8b573c4b 100644
--- a/components/private_ai/client.h
+++ b/components/private_ai/client.h
@@ -17,12 +17,9 @@
 #include "components/private_ai/proto/private_ai.pb.h"
 #include "components/private_ai/status_code.h"
 #include "components/version_info/channel.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "url/gurl.h"
 
-namespace network::mojom {
-class NetworkContext;
-}  // namespace network::mojom
-
 namespace private_ai {
 
 class PrivateAiLogger;
@@ -53,7 +50,8 @@
   // `api_key`: The API key for the PrivateAI service.
   // `proxy_url_string`: Optional URL for the proxy server.
   // `use_token_attestation`: Whether to use token attestation.
-  // `network_context`: The network context to use for connections.
+  // `network_context_getter`: Callback to get the network context to use for
+  // connections.
   // `token_manager`: Required if `use_token_attestation` is true.
   // `logger`: The logger for the client.
   // `oak_session_driver`: Interface for platform-specific capabilities related
@@ -66,7 +64,7 @@
       const std::string& api_key,
       const std::string& proxy_url_string,
       bool use_token_attestation,
-      network::mojom::NetworkContext* network_context,
+      network::NetworkContextGetter network_context_getter,
       phosphor::TokenManager* token_manager,
       PrivateAiLogger* logger,
       PrivateAiOakSessionDriver* oak_session_driver,
diff --git a/components/private_ai/connection_factory_impl.cc b/components/private_ai/connection_factory_impl.cc
index d2d94c0..68a71b84 100644
--- a/components/private_ai/connection_factory_impl.cc
+++ b/components/private_ai/connection_factory_impl.cc
@@ -24,7 +24,7 @@
 #include "components/private_ai/secure_channel.h"
 #include "components/private_ai/secure_channel_impl.h"
 #include "net/base/url_util.h"
-#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/cpp/network_context_getter.h"
 
 namespace private_ai {
 
@@ -46,13 +46,13 @@
     base::RepeatingCallback<void(StatusCode)> on_disconnect,
     base::OnceClosure on_established,
     version_info::Channel channel,
-    network::mojom::NetworkContext* network_context) {
+    network::NetworkContextGetter network_context_getter) {
   std::unique_ptr<SecureChannel::Factory> secure_channel_factory;
   if (secure_channel_override) {
     secure_channel_factory = secure_channel_override.Run();
   } else {
     secure_channel_factory = std::make_unique<SecureChannelImpl::FactoryImpl>(
-        url, network_context, logger, oak_session_driver);
+        url, std::move(network_context_getter), logger, oak_session_driver);
   }
 
   std::unique_ptr<Connection> connection = std::make_unique<ConnectionBasic>(
@@ -72,13 +72,13 @@
 
 ConnectionFactoryImpl::ConnectionFactoryImpl(
     const GURL& url,
-    network::mojom::NetworkContext* network_context,
+    network::NetworkContextGetter network_context_getter,
     PrivateAiLogger* logger,
     PrivateAiOakSessionDriver* oak_session_driver,
     PrivateAiNetworkDriver* network_driver,
     version_info::Channel channel)
     : url_(url),
-      network_context_(network_context),
+      network_context_getter_(std::move(network_context_getter)),
       logger_(logger),
       oak_session_driver_(oak_session_driver),
       network_driver_(network_driver),
@@ -114,11 +114,11 @@
   if (!proxy_url_.is_valid()) {
     logger_->LogInfo(FROM_HERE,
                      "Creating connection to Private AI server (direct).");
-    CHECK(network_context_);
+    CHECK(network_context_getter_);
     connection = CreateConnectionStack(
         url_, feature_name, logger_, token_manager_, secure_channel_override_,
         oak_session_driver_, on_disconnect, std::move(on_established), channel_,
-        network_context_);
+        network_context_getter_);
   } else {
     logger_->LogInfo(FROM_HERE,
                      "Creating connection to Private AI server via proxy: " +
diff --git a/components/private_ai/connection_factory_impl.h b/components/private_ai/connection_factory_impl.h
index 8c0ca096..d2e789db 100644
--- a/components/private_ai/connection_factory_impl.h
+++ b/components/private_ai/connection_factory_impl.h
@@ -15,12 +15,9 @@
 #include "components/private_ai/private_ai_oak_session_driver.h"
 #include "components/private_ai/secure_channel.h"
 #include "components/version_info/channel.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "url/gurl.h"
 
-namespace network::mojom {
-class NetworkContext;
-}  // namespace network::mojom
-
 namespace private_ai {
 
 class PrivateAiLogger;
@@ -35,7 +32,7 @@
       base::RepeatingCallback<std::unique_ptr<SecureChannel::Factory>()>;
 
   ConnectionFactoryImpl(const GURL& url,
-                        network::mojom::NetworkContext* network_context,
+                        network::NetworkContextGetter network_context_getter,
                         PrivateAiLogger* logger,
                         PrivateAiOakSessionDriver* oak_session_driver,
                         PrivateAiNetworkDriver* network_driver,
@@ -60,7 +57,7 @@
 
  private:
   const GURL url_;
-  const raw_ptr<network::mojom::NetworkContext> network_context_;
+  const network::NetworkContextGetter network_context_getter_;
   const raw_ptr<PrivateAiLogger> logger_;
   const raw_ptr<PrivateAiOakSessionDriver> oak_session_driver_;
   const raw_ptr<PrivateAiNetworkDriver> network_driver_;
diff --git a/components/private_ai/connection_proxy.cc b/components/private_ai/connection_proxy.cc
index dd98c3dc..df49fe7 100644
--- a/components/private_ai/connection_proxy.cc
+++ b/components/private_ai/connection_proxy.cc
@@ -8,6 +8,7 @@
 
 #include "base/base64url.h"
 #include "base/check.h"
+#include "base/functional/bind.h"
 #include "base/logging.h"
 #include "base/strings/strcat.h"
 #include "base/strings/string_util.h"
@@ -190,7 +191,13 @@
       proxied_context_.BindNewPipeAndPassReceiver(), std::move(context_params));
 
   inner_connection_ =
-      std::move(inner_connection_factory_).Run(proxied_context_.get());
+      std::move(inner_connection_factory_)
+          .Run(base::BindRepeating(
+              [](base::WeakPtr<ConnectionProxy> self)
+                  -> network::mojom::NetworkContext* {
+                return self ? self->proxied_context_.get() : nullptr;
+              },
+              weak_factory_.GetWeakPtr()));
 
   for (auto& pending : pending_requests_) {
     inner_connection_->Send(std::move(pending.request), pending.timeout,
diff --git a/components/private_ai/connection_proxy.h b/components/private_ai/connection_proxy.h
index d79c49b..354ddb1 100644
--- a/components/private_ai/connection_proxy.h
+++ b/components/private_ai/connection_proxy.h
@@ -18,6 +18,7 @@
 #include "components/private_ai/private_ai_network_driver.h"
 #include "components/private_ai/proto/private_ai.pb.h"
 #include "mojo/public/cpp/bindings/remote.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "services/network/public/mojom/network_context.mojom.h"
 #include "url/gurl.h"
 
@@ -45,9 +46,9 @@
 // the inner connection using the provided `InnerConnectionFactory`.
 class ConnectionProxy : public Connection {
  public:
-  // Factory for creating the inner connection given a NetworkContext.
+  // Factory for creating the inner connection given a NetworkContextGetter.
   using InnerConnectionFactory = base::OnceCallback<std::unique_ptr<Connection>(
-      network::mojom::NetworkContext*)>;
+      network::NetworkContextGetter)>;
 
   ConnectionProxy(const GURL& proxy_url,
                   PrivateAiLogger* logger,
diff --git a/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.cc b/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.cc
index 4ae35f11..f40b42d3 100644
--- a/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.cc
+++ b/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.cc
@@ -35,7 +35,7 @@
 #include "components/private_ai/proto/private_ai.pb.h"
 #include "components/version_info/channel.h"
 #include "mojo/public/cpp/bindings/pending_receiver.h"
-#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/cpp/network_context_getter.h"
 
 namespace private_ai {
 
@@ -79,7 +79,7 @@
 
 PrivateAiInternalsPageHandler::PrivateAiInternalsPageHandler(
     phosphor::TokenManager* token_manager,
-    network::mojom::NetworkContext* network_context,
+    network::NetworkContextGetter network_context_getter,
     Client* private_ai_client,
     PrivateAiLogger* private_ai_logger,
     PrivateAiOakSessionDriver* oak_session_driver,
@@ -92,7 +92,7 @@
       private_ai_logger_(private_ai_logger),
       oak_session_driver_(oak_session_driver),
       network_driver_(network_driver),
-      network_context_(network_context),
+      network_context_getter_(std::move(network_context_getter)),
       channel_(channel),
       receiver_(this, std::move(receiver)) {
   CHECK(private_ai_logger_);
@@ -121,7 +121,7 @@
 
   webui_client_ =
       Client::Create(url, effective_api_key, proxy_url, use_token_attestation,
-                     network_context_, token_manager_, &webui_logger_,
+                     network_context_getter_, token_manager_, &webui_logger_,
                      oak_session_driver_, network_driver_, channel_);
   std::move(callback).Run();
 }
diff --git a/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.h b/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.h
index 1d82a22..d713702 100644
--- a/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.h
+++ b/components/private_ai/private_ai_internals/webui/private_ai_internals_page_handler.h
@@ -15,10 +15,7 @@
 #include "components/version_info/channel.h"
 #include "mojo/public/cpp/bindings/receiver.h"
 #include "mojo/public/cpp/bindings/remote.h"
-
-namespace network::mojom {
-class NetworkContext;
-}  // namespace network::mojom
+#include "services/network/public/cpp/network_context_getter.h"
 
 namespace private_ai {
 
@@ -34,7 +31,7 @@
  public:
   explicit PrivateAiInternalsPageHandler(
       phosphor::TokenManager* token_manager,
-      network::mojom::NetworkContext* network_context,
+      network::NetworkContextGetter network_context_getter,
       Client* private_ai_client,
       PrivateAiLogger* private_ai_logger,
       PrivateAiOakSessionDriver* oak_session_driver,
@@ -93,7 +90,7 @@
   raw_ptr<PrivateAiOakSessionDriver> oak_session_driver_;
   raw_ptr<PrivateAiNetworkDriver> network_driver_;
   std::unique_ptr<Client> webui_client_;
-  raw_ptr<network::mojom::NetworkContext> network_context_;
+  network::NetworkContextGetter network_context_getter_;
   version_info::Channel channel_;
   mojo::Receiver<private_ai_internals::mojom::PrivateAiInternalsPageHandler>
       receiver_;
diff --git a/components/private_ai/private_ai_service.cc b/components/private_ai/private_ai_service.cc
index 196f197..e5f1835a 100644
--- a/components/private_ai/private_ai_service.cc
... (file diff truncated)

diff --git a/chrome/browser/private_ai/connection_factory_impl_browsertest.cc b/chrome/browser/private_ai/connection_factory_impl_browsertest.cc
index 8ebab9c..3d640f1 100644
--- a/chrome/browser/private_ai/connection_factory_impl_browsertest.cc
+++ b/chrome/browser/private_ai/connection_factory_impl_browsertest.cc
@@ -8,7 +8,6 @@
 #include "base/test/gtest_util.h"
 #include "base/test/scoped_feature_list.h"
 #include "chrome/browser/private_ai/private_ai_service_factory.h"
-#include "chrome/browser/profiles/profile.h"
 #include "chrome/browser/ui/omnibox/omnibox_next_features.h"
 #include "chrome/common/channel_info.h"
 #include "chrome/test/base/chrome_test_utils.h"
@@ -19,9 +18,8 @@
 #include "components/private_ai/private_ai_service.h"
 #include "components/private_ai/testing/fake_private_ai_network_driver.h"
 #include "components/private_ai/testing/fake_private_ai_oak_session_driver.h"
-#include "content/public/browser/storage_partition.h"
 #include "content/public/test/browser_test.h"
-#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "url/gurl.h"
 
@@ -44,10 +42,9 @@
   ~ConnectionFactoryImplBrowserTest() override = default;
 
  protected:
-  network::mojom::NetworkContext* GetNetworkContext() {
-    return chrome_test_utils::GetProfile(this)
-        ->GetDefaultStoragePartition()
-        ->GetNetworkContext();
+  network::NetworkContextGetter GetNetworkContextGetter() {
+    return PrivateAiServiceFactory::CreateNetworkContextGetter(
+        chrome_test_utils::GetProfile(this));
   }
 
   phosphor::TokenManager* GetTokenManager() {
@@ -77,7 +74,7 @@
                        CreateConnectionWithoutToken) {
   GURL url("wss://private-ai.googleapis.com?key=test_api_key");
 
-  ConnectionFactoryImpl factory(url, GetNetworkContext(), GetLogger(),
+  ConnectionFactoryImpl factory(url, GetNetworkContextGetter(), GetLogger(),
                                 GetOakSessionDriver(), GetNetworkDriver(),
                                 chrome::GetChannel());
 
@@ -98,7 +95,7 @@
                        MAYBE_FactoryCtorFailsWithoutApiKey) {
   GURL url("wss://private-ai.googleapis.com");
   EXPECT_CHECK_DEATH(ConnectionFactoryImpl(
-      url, GetNetworkContext(), GetLogger(), GetOakSessionDriver(),
+      url, GetNetworkContextGetter(), GetLogger(), GetOakSessionDriver(),
       GetNetworkDriver(), chrome::GetChannel()));
 }
 
@@ -106,7 +103,7 @@
                        CreateConnectionWithToken) {
   GURL url("wss://private-ai.googleapis.com?key=test_api_key");
 
-  ConnectionFactoryImpl factory(url, GetNetworkContext(), GetLogger(),
+  ConnectionFactoryImpl factory(url, GetNetworkContextGetter(), GetLogger(),
                                 GetOakSessionDriver(), GetNetworkDriver(),
                                 chrome::GetChannel());
   factory.EnableTokenAttestation(GetTokenManager());
@@ -121,7 +118,7 @@
                        CreateConnectionWithProxyAndToken) {
   GURL url("wss://private-ai.googleapis.com?key=test_api_key");
 
-  ConnectionFactoryImpl factory(url, GetNetworkContext(), GetLogger(),
+  ConnectionFactoryImpl factory(url, GetNetworkContextGetter(), GetLogger(),
                                 GetOakSessionDriver(), GetNetworkDriver(),
                                 chrome::GetChannel());
   factory.EnableTokenAttestation(GetTokenManager());
diff --git a/chrome/browser/private_ai/private_ai_service_browsertest.cc b/chrome/browser/private_ai/private_ai_service_browsertest.cc
index 401566c4..6ededc9 100644
--- a/chrome/browser/private_ai/private_ai_service_browsertest.cc
+++ b/chrome/browser/private_ai/private_ai_service_browsertest.cc
@@ -52,7 +52,7 @@
               IdentityManagerFactory::GetForProfile(profile),
               profile->GetDefaultStoragePartition()
                   ->GetURLLoaderFactoryForBrowserProcess(),
-              profile->GetDefaultStoragePartition()->GetNetworkContext(),
+              PrivateAiServiceFactory::CreateNetworkContextGetter(profile),
               kPrivateAiUrl.Get(),
               PrivateAiService::GetApiKey(chrome::GetChannel()),
               kPrivateAiProxyServerUrl.Get(),
diff --git a/components/private_ai/client_impl_integration_test.cc b/components/private_ai/client_impl_integration_test.cc
index bba0211..26d0dec 100644
--- a/components/private_ai/client_impl_integration_test.cc
+++ b/components/private_ai/client_impl_integration_test.cc
@@ -51,8 +51,12 @@
     GURL url("wss://example.com?key=test-api-key");
 
     auto factory = std::make_unique<ConnectionFactoryImpl>(
-        url, &test_network_context_, &logger_, &oak_session_driver_,
-        &network_driver_, version_info::Channel::UNKNOWN);
+        url,
+        base::BindRepeating(
+            [](network::mojom::NetworkContext* context) { return context; },
+            base::Unretained(&test_network_context_)),
+        &logger_, &oak_session_driver_, &network_driver_,
+        version_info::Channel::UNKNOWN);
     factory_ptr_ = factory.get();
     factory->EnableTokenAttestation(&token_manager_);
     factory->SetSecureChannelFactoryForTesting(base::BindLambdaForTesting(
diff --git a/components/private_ai/connection_proxy_unittest.cc b/components/private_ai/connection_proxy_unittest.cc
index 332fdc9..d9c2cf8 100644
--- a/components/private_ai/connection_proxy_unittest.cc
+++ b/components/private_ai/connection_proxy_unittest.cc
@@ -21,7 +21,7 @@
 #include "components/private_ai/testing/fake_private_ai_network_driver.h"
 #include "components/private_ai/testing/fake_token_manager.h"
 #include "net/base/proxy_string_util.h"
-#include "services/network/public/mojom/network_context.mojom.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
@@ -47,7 +47,7 @@
   }
 
   std::unique_ptr<Connection> CreateInnerConnection(
-      network::mojom::NetworkContext* context) {
+      network::NetworkContextGetter context_getter) {
     auto connection = std::make_unique<FakeConnection>(base::DoNothing());
     inner_connection_ = connection.get();
     if (on_inner_connection_created_) {
diff --git a/components/private_ai/private_ai_service_unittest.cc b/components/private_ai/private_ai_service_unittest.cc
index 7d17ffe..c48caf96 100644
--- a/components/private_ai/private_ai_service_unittest.cc
+++ b/components/private_ai/private_ai_service_unittest.cc
@@ -9,6 +9,7 @@
 #include <vector>
 
 #include "base/base64.h"
+#include "base/functional/bind.h"
 #include "base/test/bind.h"
 #include "base/test/scoped_feature_list.h"
 #include "base/test/task_environment.h"
@@ -74,7 +75,10 @@
 
     private_ai_service_ = std::make_unique<TestPrivateAiService>(
         identity_test_env_.identity_manager(),
-        std::move(shared_url_loader_factory), &test_network_context_,
+        std::move(shared_url_loader_factory),
+        base::BindRepeating(
+            [](network::mojom::NetworkContext* context) { return context; },
+            base::Unretained(&test_network_context_)),
         "dummy.com", PrivateAiService::GetApiKey(version_info::Channel::STABLE),
         "dummy-proxy.com",
         /*use_token_attestation=*/false,
diff --git a/components/private_ai/testing/test_private_ai_service.cc b/components/private_ai/testing/test_private_ai_service.cc
index 3f90842..fe5ab304 100644
--- a/components/private_ai/testing/test_private_ai_service.cc
+++ b/components/private_ai/testing/test_private_ai_service.cc
@@ -4,8 +4,11 @@
 
 #include "components/private_ai/testing/test_private_ai_service.h"
 
+#include <utility>
+
 #include "components/private_ai/private_ai_network_driver.h"
 #include "components/private_ai/private_ai_oak_session_driver.h"
+#include "services/network/public/cpp/network_context_getter.h"
 #include "services/network/public/cpp/shared_url_loader_factory.h"
 
 namespace private_ai {
@@ -13,7 +16,7 @@
 TestPrivateAiService::TestPrivateAiService(
     signin::IdentityManager* identity_manager,
     scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
-    network::mojom::NetworkContext* network_context,
+    network::NetworkContextGetter network_context_getter,
     const std::string& url,
     const std::string& api_key,
     const std::string& proxy_url,
@@ -27,7 +30,7 @@
                        std::move(url_loader_factory),
                        std::move(network_driver),
                        std::move(oak_session_driver),
-                       network_context,
+                       std::move(network_context_getter),
                        url,
                        api_key,
                        proxy_url,
diff --git a/components/private_ai/testing/test_private_ai_service.h b/components/private_ai/testing/test_private_ai_service.h
index bc76ff24..fe209b0 100644
--- a/components/private_ai/testing/test_private_ai_service.h
+++ b/components/private_ai/testing/test_private_ai_service.h
@@ -11,15 +11,12 @@
 #include "base/memory/scoped_refptr.h"
 #include "components/private_ai/private_ai_service.h"
 #include "components/private_ai/testing/test_blind_sign_auth_factory.h"
+#include "services/network/public/cpp/network_context_getter.h"
 
 namespace network {
 class SharedURLLoaderFactory;
 }
 
-namespace network::mojom {
-class NetworkContext;
-}
-
 namespace signin {
 class IdentityManager;
 }
@@ -40,7 +37,7 @@
   TestPrivateAiService(
       signin::IdentityManager* identity_manager,
       scoped_refptr<network::SharedURLLoaderFactory> url_loader_factory,
-      network::mojom::NetworkContext* network_context,
+      network::NetworkContextGetter network_context_getter,
       const std::string& url,
       const std::string& api_key,
       const std::string& proxy_url,