Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
68 changes: 18 additions & 50 deletions src/google/adk/telemetry/google_cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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,
)

Expand Down Expand Up @@ -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(
Expand All @@ -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,
),
)
Original file line number Diff line number Diff line change
Expand Up @@ -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
2 changes: 1 addition & 1 deletion tests/unittests/telemetry/test_google_cloud.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading