Skip to content

Commit babf2a7

Browse files
caohy1988claude
andcommitted
fix: clear both _credentials and _user_credentials on pickle
Addresses review feedback: - Non-picklable user-provided credentials (e.g., compute_engine with requests.Session) broke pickle.dumps() because _user_credentials was preserved in __getstate__. Now both _credentials and _user_credentials are cleared; credentials re-resolve via ADC after unpickle. - Fork safety: _reset_runtime_state documents that user-provided credentials are kept as-is (we cannot re-create them), while ADC-resolved credentials are cleared for re-resolution. - GCS client: always pass credentials explicitly when available; GCSOffloader.__init__ always creates a client eagerly (the "lazy restoration" claim in the previous commit was incorrect — the pre-credentials code had the same eager behavior). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 6199e11 commit babf2a7

1 file changed

Lines changed: 17 additions & 13 deletions

File tree

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

Lines changed: 17 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -2201,17 +2201,18 @@ async def _lazy_setup(self, **kwargs) -> None:
22012201

22022202
self.offloader = None
22032203
if self.config.gcs_bucket_name:
2204-
gcs_client = None
2204+
# GCSOffloader always creates a storage.Client eagerly
2205+
# (line 1329: storage_client or storage.Client(...)).
2206+
# Pass credentials so it uses the same auth as the other
2207+
# clients; omit when None to let it use ADC.
2208+
gcs_kwargs = {"project": self.project_id}
22052209
if self._credentials is not None:
2206-
gcs_client = storage.Client(
2207-
project=self.project_id,
2208-
credentials=self._credentials,
2209-
)
2210+
gcs_kwargs["credentials"] = self._credentials
22102211
self.offloader = GCSOffloader(
22112212
self.project_id,
22122213
self.config.gcs_bucket_name,
22132214
self._executor,
2214-
storage_client=gcs_client,
2215+
storage_client=storage.Client(**gcs_kwargs),
22152216
)
22162217

22172218
self.parser = HybridContentParser(
@@ -2547,8 +2548,10 @@ def __getstate__(self):
25472548
state["_init_pid"] = 0
25482549
# Credential objects may hold non-picklable transport state
25492550
# (e.g., requests.Session in compute_engine.Credentials).
2550-
# Clear and re-resolve from _user_credentials on unpickle.
2551+
# Clear both so pickle succeeds regardless of credential type.
2552+
# After unpickle, credentials are re-resolved via ADC.
25512553
state["_credentials"] = None
2554+
state["_user_credentials"] = None
25522555
return state
25532556

25542557
def __setstate__(self, state):
@@ -2557,11 +2560,10 @@ def __setstate__(self, state):
25572560
# code versions so _ensure_started does not raise AttributeError.
25582561
state.setdefault("_init_pid", 0)
25592562
state.setdefault("_user_credentials", None)
2560-
# Restore _credentials from _user_credentials (if user provided
2561-
# them); ADC-resolved credentials will be re-resolved on next
2563+
state.setdefault("_credentials", None)
2564+
# Both _credentials and _user_credentials are cleared during
2565+
# pickle. Credentials will be re-resolved via ADC on the next
25622566
# _create_loop_state call.
2563-
if state.get("_credentials") is None:
2564-
state["_credentials"] = state.get("_user_credentials")
25652567
self.__dict__.update(state)
25662568

25672569
def _reset_runtime_state(self) -> None:
@@ -2616,8 +2618,10 @@ def _reset_runtime_state(self) -> None:
26162618
self._startup_error = None
26172619
self._is_shutting_down = False
26182620
self._init_pid = os.getpid()
2619-
# Credentials may hold stale HTTP transport state after fork.
2620-
# Re-resolve from _user_credentials on next _create_loop_state.
2621+
# For ADC-resolved credentials, clear so they are re-resolved
2622+
# in the child process. For user-provided credentials, keep
2623+
# the original object — we cannot re-create it. The user is
2624+
# responsible for providing fork-safe credentials if needed.
26212625
self._credentials = self._user_credentials
26222626

26232627
async def __aenter__(self) -> BigQueryAgentAnalyticsPlugin:

0 commit comments

Comments
 (0)