Skip to content
Open
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
78 changes: 76 additions & 2 deletions src/google/adk/cli/trigger_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
- Semaphore-based concurrency control to stay within LLM model quota
- Automatic retry with exponential backoff on 429 / RESOURCE_EXHAUSTED
- Transient error detection to signal upstream services to retry
- Delivery identity in session state so tools can recognize redeliveries
"""

from __future__ import annotations
Expand Down Expand Up @@ -82,6 +83,42 @@
)
"""Maximum delay in seconds for exponential backoff."""

TRIGGER_DELIVERY_STATE_KEY = "trigger_delivery"
"""Session state key holding the identity of the delivery that started a run.

Pub/Sub and Eventarc redeliver a message whenever the endpoint does not
acknowledge it, and every delivery runs the agent in a new session with fresh
state, so nothing an earlier attempt wrote to session state is visible to the
redelivery. A tool with side effects should instead derive an idempotency key
from this value and deduplicate where the side effect happens:

- Prefer the provider's idempotency key (for example a payment API's). Only
the provider knows whether the effect happened.
- If you use your own store instead, reserve the key atomically before the
effect, and plan for reconciling an ambiguous attempt, such as a call that
timed out after the provider committed. Checking the store, acting and
then recording the key repeats the effect in that case, and when a
redelivery overlaps a run still in flight.
- If the provider answers that the key is still in use by another attempt
(for example an HTTP 409), raise from the tool instead of returning the
error to the model, so the message is nacked and redelivered.

The value is a dict with:

- ``source``: ``"pubsub"`` or ``"eventarc"``.
- ``id``: the Pub/Sub ``messageId`` or the CloudEvents ``id``, which stay
the same across redeliveries of one message. ``None`` if the request did
not carry one.
- ``subscription`` and ``publish_time`` for Pub/Sub, or ``event_source``
and ``type`` for Eventarc.

``id`` alone is not globally unique: a Pub/Sub ``messageId`` is unique only
within its topic, and every subscription on the topic receives the message,
while a CloudEvents ``id`` is unique only within its source. Build the key
from ``subscription`` and ``id`` for Pub/Sub, and from ``event_source`` and
``id`` for Eventarc.
"""


# ---------------------------------------------------------------------------
# Transient error detection
Expand Down Expand Up @@ -341,6 +378,7 @@ async def _run_agent(
user_id: str,
message_text: str,
session_id: str,
session_state: Optional[dict[str, Any]] = None,
) -> list[Event]:
"""Run the agent with an auto-created ephemeral session.

Expand All @@ -352,6 +390,7 @@ async def _run_agent(
user_id: Identifier for observability (derived from trigger metadata).
message_text: The text input to send to the agent.
session_id: The session ID to use.
session_state: Initial state for the session if it has to be created.

Returns:
List of events produced by the agent invocation.
Expand All @@ -369,6 +408,7 @@ async def _run_agent(
session = await self._server.session_service.create_session(
app_name=app_name,
user_id=user_id,
state=session_state,
session_id=session_id,
)

Expand Down Expand Up @@ -396,6 +436,7 @@ async def _run_agent_with_retry(
app_name: str,
user_id: str,
message_text: str,
delivery: Optional[dict[str, Any]] = None,
) -> list[Event]:
"""Run the agent with retry on transient errors.

Expand All @@ -407,6 +448,8 @@ async def _run_agent_with_retry(
app_name: The target application / agent name.
user_id: Identifier for observability.
message_text: The text input to send to the agent.
delivery: Identity of the upstream delivery, stored in session state
under ``TRIGGER_DELIVERY_STATE_KEY``.

Returns:
List of events produced by the agent invocation.
Expand All @@ -417,6 +460,9 @@ async def _run_agent_with_retry(
"""
last_error: Optional[Exception] = None
session_id = str(uuid.uuid4())
session_state = (
{TRIGGER_DELIVERY_STATE_KEY: delivery} if delivery is not None else None
)

for attempt in range(self._max_retries + 1):
try:
Expand All @@ -425,6 +471,7 @@ async def _run_agent_with_retry(
user_id=user_id,
message_text=message_text,
session_id=session_id,
session_state=session_state,
)
except Exception as e:
if not _is_transient_error(e):
Expand Down Expand Up @@ -485,6 +532,9 @@ async def _verify_auth(request: Request) -> None:
"Processes a message from a Pub/Sub push subscription."
" Returns 200 on success; errors trigger Pub/Sub retry."
" Includes automatic retry with backoff on 429 errors."
" The messageId is stored in session state under"
f" '{TRIGGER_DELIVERY_STATE_KEY}' so tools can detect"
" redeliveries."
),
dependencies=auth_dependencies,
)
Expand Down Expand Up @@ -526,6 +576,12 @@ async def trigger_pubsub(
app_name=app_name,
user_id=user_id,
message_text=message_text,
delivery={
"source": "pubsub",
"id": req.message.messageId,
"subscription": req.subscription,
"publish_time": req.message.publishTime,
},
)
except TransientError as te:
logger.exception("Pub/Sub: transient error after retries: %s", te)
Expand Down Expand Up @@ -553,6 +609,9 @@ async def trigger_pubsub(
"Processes a CloudEvent delivered by Eventarc."
" Returns 200 on success; errors trigger Eventarc retry."
" Includes automatic retry with backoff on 429 errors."
" The event id is stored in session state under"
f" '{TRIGGER_DELIVERY_STATE_KEY}' so tools can detect"
" redeliveries."
),
dependencies=auth_dependencies,
)
Expand All @@ -564,11 +623,18 @@ async def trigger_eventarc(
default="eventarc-caller",
)

event_type = req.type or request.headers.get("ce-type")
event_id = req.id or request.headers.get("ce-id")
if not event_id and req.message:
# A Pub/Sub-wrapped event may arrive without a CloudEvents id; its
# messageId is just as stable across redeliveries.
event_id = req.message.messageId

logger.info(
"Eventarc trigger: source=%s, type=%s, id=%s",
user_id,
req.type or request.headers.get("ce-type"),
req.id or request.headers.get("ce-id"),
event_type,
event_id,
)

# Extract message text — support both structured and binary modes.
Expand Down Expand Up @@ -644,6 +710,14 @@ async def trigger_eventarc(
app_name=app_name,
user_id=user_id,
message_text=message_text,
delivery={
"source": "eventarc",
"id": event_id,
"event_source": (
req.source or request.headers.get("ce-source")
),
"type": event_type,
},
)
except TransientError as te:
logger.exception("Eventarc: transient error after retries: %s", te)
Expand Down
Loading
Loading