From f782cc7e4dcd1be7f3560c74e0af25edea858316 Mon Sep 17 00:00:00 2001
From: Walt Wells <walt@khanacademy.org>
Date: Wed, 4 Jun 2025 13:02:26 -0400
Subject: [PATCH] update GCSBucketBase to handle GSM project ID if passed
 (#11409)

---
 .../gcs_bucket/gcs_bucket_base.py             | 10 ++-
 .../gcs_bucket/test_gcs_bucket_base.py        | 82 +++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py

diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py
index 0ce845ecb2..2612face05 100644
--- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py
+++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py
@@ -66,11 +66,19 @@ class GCSBucketBase(CustomBatchLogger):
         return headers
 
     def sync_construct_request_headers(self) -> Dict[str, str]:
+        """
+        Construct request headers for GCS API calls
+        """
         from litellm import vertex_chat_completion
 
+        # Get project_id from environment if available, otherwise None
+        # This helps support use of this library to auth to pull secrets 
+        # from Secret Manager.
+        project_id = os.getenv("GOOGLE_SECRET_MANAGER_PROJECT_ID")
+        
         _auth_header, vertex_project = vertex_chat_completion._ensure_access_token(
             credentials=self.path_service_account_json,
-            project_id=None,
+            project_id=project_id,
             custom_llm_provider="vertex_ai",
         )
 
diff --git a/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py
new file mode 100644
index 0000000000..80925d02f1
--- /dev/null
+++ b/tests/test_litellm/integrations/gcs_bucket/test_gcs_bucket_base.py
@@ -0,0 +1,82 @@
+import os
+from unittest.mock import patch
+from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase
+
+
+class TestGCSBucketBase:
+    def test_construct_request_headers_with_project_id(self):
+        """Test that construct_request_headers correctly uses project_id if passed from env"""
+        test_project_id = "test-project"
+        os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"] = test_project_id
+
+        try:
+            # Create handler
+            handler = GCSBucketBase(bucket_name="test-bucket")
+
+            # Mock the Vertex AI auth calls
+            mock_auth_header = "mock-auth-header"
+            mock_token = "mock-token"
+
+            with patch(
+                "litellm.vertex_chat_completion._ensure_access_token"
+            ) as mock_ensure_token, patch(
+                "litellm.vertex_chat_completion._get_token_and_url"
+            ) as mock_get_token:
+                mock_ensure_token.return_value = (mock_auth_header, test_project_id)
+                mock_get_token.return_value = (mock_token, "mock-url")
+
+                # Call construct_request_headers
+                headers = handler.sync_construct_request_headers()
+
+                # Verify headers
+                assert headers == {
+                    "Authorization": f"Bearer {mock_token}",
+                    "Content-Type": "application/json",
+                }
+
+                # Verify _ensure_access_token was called with correct project_id
+                mock_ensure_token.assert_called_once_with(
+                    credentials=None,  # No service account in this test
+                    project_id=test_project_id,  # Should use project_id from env
+                    custom_llm_provider="vertex_ai",
+                )
+        finally:
+            # Clean up environment variable
+            del os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"]
+
+    def test_construct_request_headers_without_project_id(self):
+        """Test that construct_request_headers works when no project_id is in env"""
+        # Ensure environment variable is not set
+        if "GOOGLE_SECRET_MANAGER_PROJECT_ID" in os.environ:
+            del os.environ["GOOGLE_SECRET_MANAGER_PROJECT_ID"]
+
+        # Create handler
+        handler = GCSBucketBase(bucket_name="test-bucket")
+
+        # Mock the Vertex AI auth calls
+        mock_auth_header = "mock-auth-header"
+        mock_token = "mock-token"
+
+        with patch(
+            "litellm.vertex_chat_completion._ensure_access_token"
+        ) as mock_ensure_token, patch(
+            "litellm.vertex_chat_completion._get_token_and_url"
+        ) as mock_get_token:
+            mock_ensure_token.return_value = (mock_auth_header, None)
+            mock_get_token.return_value = (mock_token, "mock-url")
+
+            # Call construct_request_headers
+            headers = handler.sync_construct_request_headers()
+
+            # Verify headers
+            assert headers == {
+                "Authorization": f"Bearer {mock_token}",
+                "Content-Type": "application/json",
+            }
+
+            # Verify _ensure_access_token was called with project_id as None
+            mock_ensure_token.assert_called_once_with(
+                credentials=None,  # No service account in this test
+                project_id=None,  # Should be None when no env var is set
+                custom_llm_provider="vertex_ai",
+            )