diff --git a/src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py b/src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py index 9d8e702f703..97716bb785c 100644 --- a/src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py +++ b/src/google/adk/integrations/agent_identity/_iam_connector_credentials_provider.py @@ -238,7 +238,7 @@ async def get_auth_credential( logger.debug("Auth credential obtained immediately.") return _construct_auth_credential(response) - if metadata and metadata.consent_pending: + if metadata is not None and "consent_pending" in metadata: # Get 2-legged OAuth token. Allow enough time for token exchange. try: operation = await self._poll_credentials( diff --git a/src/google/adk/telemetry/google_cloud.py b/src/google/adk/telemetry/google_cloud.py index a224893fe98..0f1e12e90f5 100644 --- a/src/google/adk/telemetry/google_cloud.py +++ b/src/google/adk/telemetry/google_cloud.py @@ -119,7 +119,6 @@ def get_gcp_exporters( if enable_cloud_logging: exporter = _get_gcp_logs_exporter( project_id=project_id, - credentials=credentials, ) if exporter: log_record_processors.append(exporter) @@ -187,11 +186,9 @@ def _get_gcp_metrics_exporter(project_id: str) -> MetricReader: def _get_gcp_logs_exporter( project_id: str, - credentials: Optional["Credentials"] = None, ) -> LogRecordProcessor: if os.getenv("GOOGLE_CLOUD_AGENT_ENGINE_ID"): return _get_agent_engine_logs_exporter( - credentials=credentials, project_id=project_id, ) @@ -338,18 +335,14 @@ def _use_client_cert_effective() -> bool: def _get_agent_engine_logs_exporter( *, - credentials: "Credentials", project_id: str, ): """Configures logging for Agent Engine. Args: - credentials: Credentials to use for export calls. project_id: Project to which to write logs. """ try: - from google.cloud.logging_v2.services import logging_service_v2 - from google.cloud.logging_v2.services.logging_service_v2.transports import grpc from opentelemetry.exporter import cloud_logging except (ImportError, AttributeError): logger.warning( @@ -363,46 +356,21 @@ def _get_agent_engine_logs_exporter( ) return - if "gen_ai_latest_experimental" in os.getenv( - "OTEL_SEMCONV_STABILITY_OPT_IN", "" - ).split(","): - # Specify credentials to avoid expensive call to `google.auth.default()` - channel = grpc.LoggingServiceV2GrpcTransport.create_channel( - credentials=credentials, - # pylint: disable-next=protected-access - options=cloud_logging._OPTIONS, - ) - return BatchLogRecordProcessor( - cloud_logging.CloudLoggingExporter( - client=logging_service_v2.LoggingServiceV2Client( - transport=grpc.LoggingServiceV2GrpcTransport( - credentials=credentials, - channel=channel, - ), - ), - project_id=project_id, - default_log_name=os.getenv( - "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine" - ), - ), - ) - else: - - class _SimpleLogRecordProcessor(SimpleLogRecordProcessor): - - def force_flush( - self, timeout_millis: int = 30000 - ) -> bool: # pylint: disable=no-self-use - _ = sys.stdout.flush() - _ = sys.stderr.flush() - return super().force_flush() - - return _SimpleLogRecordProcessor( - cloud_logging.CloudLoggingExporter( - project_id=project_id, - default_log_name=os.getenv( - "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine" - ), - structured_json_file=sys.stdout, - ), - ) + class _SimpleLogRecordProcessor(SimpleLogRecordProcessor): + + def force_flush( + self, timeout_millis: int = 30000 + ) -> bool: # pylint: disable=no-self-use + _ = sys.stdout.flush() + _ = sys.stderr.flush() + return super().force_flush() + + return _SimpleLogRecordProcessor( + cloud_logging.CloudLoggingExporter( + project_id=project_id, + default_log_name=os.getenv( + "GCP_DEFAULT_LOG_NAME", "adk-on-agent-engine" + ), + structured_json_file=sys.stdout, + ), + ) diff --git a/tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py b/tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py index 1ecadf6e2e7..2eb279dcdfe 100644 --- a/tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py +++ b/tests/unittests/integrations/agent_identity/test_iam_connector_credentials_provider.py @@ -469,3 +469,42 @@ async def test_get_auth_credential_raises_error_if_consent_canceled( RuntimeError, match="Failed to retrieve consent based credential." ): await provider.get_auth_credential(auth_scheme, context) + + +async def test_get_auth_credential_handles_consent_pending_state_correctly( + mock_client, + auth_scheme, + context, + provider, +): + """Test get_auth_credential enters the polling loop when consent is pending.""" + # 1. Setup the first retrieve_credentials call to return a pending Operation + # containing consent_pending metadata. + meta_pb = RetrieveCredentialsMetadata.pb()() + meta_pb.consent_pending.SetInParent() + + op_pending = Operation(done=False) + op_pending.metadata.value = meta_pb.SerializeToString() + + # 2. Setup the second retrieve_credentials call (the poll) to return success. + op_success = Operation(done=True) + resp = RetrieveCredentialsResponse( + header="Authorization: Bearer", token="valid-token" + ) + op_success.response.value = RetrieveCredentialsResponse.serialize(resp) + + # Configure mock client to return pending first, then success + mock_client.retrieve_credentials.side_effect = [ + Mock(operation=op_pending), + Mock(operation=op_success), + ] + + # 3. Call the provider + credential = await provider.get_auth_credential(auth_scheme, context) + + # 4. Verify that it polled and successfully returned the token + assert credential is not None + assert credential.http.credentials.token == "valid-token" + + # Verify that retrieve_credentials was called twice (initial + 1 poll) + assert mock_client.retrieve_credentials.call_count == 2 diff --git a/tests/unittests/telemetry/test_google_cloud.py b/tests/unittests/telemetry/test_google_cloud.py index 8b57ac9dd11..7559f9f2011 100644 --- a/tests/unittests/telemetry/test_google_cloud.py +++ b/tests/unittests/telemetry/test_google_cloud.py @@ -62,7 +62,7 @@ def test_get_gcp_exporters( ) monkeypatch.setattr( "google.adk.telemetry.google_cloud._get_gcp_logs_exporter", - lambda project_id, credentials: mock.MagicMock(), + lambda project_id: mock.MagicMock(), ) # Act.