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
101 changes: 38 additions & 63 deletions src/bedrock_agentcore/payments/integrations/strands/plugin.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,11 +163,6 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None:
)
return

# Check if payment retry limit has been reached
if self._check_payment_retry_limit(event):
logger.warning("Payment processing retry limit has been reached. Processing skipped.")
return

# Check if response is a 402 Payment Required
if not hasattr(event, "result") or event.result is None:
return
Expand All @@ -192,9 +187,6 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None:

logger.info("Detected 402 Payment Required response from tool: %s", event.tool_use.get("name", "unknown"))

# Increment retry count in invocation state
self._increment_payment_retry_count(event)

# Build payment_required_request dict using handler methods
headers = handler.extract_headers(event.result)
body = handler.extract_body(event.result)
Expand All @@ -204,20 +196,28 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None:
"body": body or {},
}

# If we already retried with payment credentials and still got a 402,
# this is a post-payment failure (e.g., insufficient balance, invalid signature).
# Propagate as an interrupt instead of retrying again to avoid infinite loops.
if self._is_post_payment_failure(event, body):
# If we previously signed successfully and still got a 402, the server
# rejected the payment for a non-retryable reason (e.g., insufficient balance).
# Do not retry — store failure state so the agent is notified via interrupt.
if self._has_successful_signing(event):
error_msg = body.get("error", "unknown error") if body and isinstance(body, dict) else "unknown error"
logger.warning(
"Received 402 after payment retry for tool %s — treating as payment failure",
"Received 402 after successful signing for tool %s — post-payment failure: %s",
event.tool_use.get("name", "unknown"),
error_msg,
)
error_msg = self._extract_payment_error_message(body)
self._store_payment_failure_state(event, PaymentError(f"Payment failed after retry: {error_msg}"))
self._store_payment_failure_state(event, PaymentError(f"Payment rejected after signing: {error_msg}"))
return

# Check if signing retry limit has been reached
if self._check_payment_retry_limit(event):
logger.warning("Payment signing retry limit reached. Processing skipped.")
return

# Increment before attempt so limit is enforced even on exception
self._increment_payment_retry_count(event)

# Validate tool input before processing payment
tool_input = event.tool_use.get("input", {})
if not handler.validate_tool_input(tool_input):
logger.error("Tool input validation failed, cannot apply payment header")
self._store_payment_failure_state(event, Exception("Tool input validation failed"))
Expand All @@ -232,10 +232,11 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None:
self._store_payment_failure_state(event, Exception("Failed to apply payment header"))
return

# Mark that signing succeeded for this tool use — if we get another 402
# after this retry, we know it's a server-side rejection, not a signing failure.
self._mark_successful_signing(event)

# Set retry flag to re-execute the tool with payment credentials.
# Do NOT reset the payment retry counter here — it must persist across
# retries so that _is_post_payment_failure and _check_payment_retry_limit
# can detect repeated 402s and break the loop.
event.retry = True
self._reset_interrupt_retry_count(event)
logger.info("Set retry flag to re-execute tool with payment credentials")
Expand Down Expand Up @@ -295,63 +296,37 @@ def _increment_payment_retry_count(self, event: AfterToolCallEvent) -> None:
"Payment retry attempt %d/%d for tool use %s", retry_count + 1, self.MAX_PAYMENT_RETRIES, tool_use_id
)

def _is_post_payment_failure(self, event: AfterToolCallEvent, body: Optional[Dict[str, Any]]) -> bool:
"""Check if this 402 response is a failure after we already retried with payment credentials.

A post-payment failure occurs when:
1. We already sent a payment header (retry count > 0 before this increment), AND
2. The 402 response body contains an error that is NOT the initial "payment required"
(e.g., "invalid_exact_evm_insufficient_balance", "payment_rejected", etc.)

This prevents infinite loops where the plugin keeps signing and retrying
against a server that keeps rejecting the payment for non-retryable reasons.
def _has_successful_signing(self, event: AfterToolCallEvent) -> bool:
"""Check if we previously signed a payment successfully for this tool use.

Args:
event: The after tool call event
body: The extracted response body (may be None)

Returns:
True if this is a post-payment failure that should be propagated as an interrupt
True if signing was previously successful (meaning this 402 is a server-side rejection)
"""
tool_use_id = event.tool_use.get("toolUseId", "unknown")
payment_retry_key = f"payment_retry_count_{tool_use_id}"
# retry count was already incremented before this check, so > 1 means
# we already attempted at least one payment retry
retry_count = event.invocation_state.get(payment_retry_key, 0)

if retry_count <= 1:
return False
signed_key = f"payment_signed_{tool_use_id}"
return event.invocation_state.get(signed_key, False)

# If the body contains an error field that is NOT the initial "payment required",
# this is a post-payment failure
if body and isinstance(body, dict):
error = body.get("error", "")
if isinstance(error, str) and error.lower() not in ("", "payment required"):
logger.info(
"Post-payment failure detected for tool %s: error=%s (retry_count=%d)",
tool_use_id,
error,
retry_count,
)
return True
def _mark_successful_signing(self, event: AfterToolCallEvent) -> None:
"""Mark that signing succeeded for this tool use.

return False
Called after generate_payment_header and apply_payment_header both succeed,
right before setting event.retry. If a subsequent 402 is received,
_has_successful_signing will return True indicating the failure is server-side.

@staticmethod
def _extract_payment_error_message(body: Optional[Dict[str, Any]]) -> str:
"""Extract a human-readable error message from a 402 response body.
Note: payment_signed_*, payment_retry_count_*, and payment_failure_* keys are
intentionally not cleared. invocation_state is scoped to a single agent
invocation and is discarded by Strands when the invocation ends, so these
per-tool-use markers do not accumulate across invocations.

Args:
body: The extracted response body (may be None)

Returns:
Error message string, or "unknown error" if not extractable
event: The after tool call event
"""
if body and isinstance(body, dict):
error = body.get("error")
if isinstance(error, str) and error:
return error
return "unknown error"
tool_use_id = event.tool_use.get("toolUseId", "unknown")
signed_key = f"payment_signed_{tool_use_id}"
event.invocation_state[signed_key] = True
Comment on lines +312 to +329
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nit: payment_signed_* flag is never cleared.

Once _mark_successful_signing sets the flag, it lives in invocation_state for the rest of the invocation. Same applies to payment_retry_count_* and payment_failure_* (pre-existing). For long-lived agent sessions with many paid tool calls, invocation_state accumulates keys without bound. Either clear the flag once the failure is delivered via before_tool_call's interrupt path, or add a comment confirming invocation_state is intentionally per-invocation and cleanup is not required.


def _store_payment_failure_state(self, event: AfterToolCallEvent, exception: Exception) -> None:
"""Store payment failure information in invocation state for agent to handle.
Expand Down
Loading
Loading