From 34f91305cb764a9b7f79d96ba6adebadb924a62f Mon Sep 17 00:00:00 2001 From: Raju Ansari Date: Thu, 21 May 2026 10:52:21 -0700 Subject: [PATCH 1/3] fix: stop retrying after successful payment signing is rejected by server Previously, when a signed payment was rejected (e.g., insufficient balance), the plugin would store failure state and trigger interrupt-based retries. This was unnecessary since a server-side rejection after successful signing is not recoverable by retrying. Changes: - Add payment_signed_{toolUseId} flag to track successful signing - After successful signing + retry, if server returns 402 again, stop immediately - Signing failures still retry up to MAX_PAYMENT_RETRIES - Move retry count increment to after successful signing attempt - Remove _is_post_payment_failure (replaced by _has_successful_signing) --- .../payments/integrations/strands/plugin.py | 82 ++--- .../integrations/strands/test_plugin.py | 183 ++++++----- .../strands/test_plugin_integration.py | 295 ++++++++++++++++++ 3 files changed, 430 insertions(+), 130 deletions(-) diff --git a/src/bedrock_agentcore/payments/integrations/strands/plugin.py b/src/bedrock_agentcore/payments/integrations/strands/plugin.py index edc8372d..3cf55ab6 100644 --- a/src/bedrock_agentcore/payments/integrations/strands/plugin.py +++ b/src/bedrock_agentcore/payments/integrations/strands/plugin.py @@ -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 @@ -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) @@ -204,18 +196,24 @@ 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). + # Stop processing — the existing 402 result is already structured for the LLM. + if self._has_successful_signing(event): 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, not retrying", event.tool_use.get("name", "unknown"), ) - error_msg = self._extract_payment_error_message(body) - self._store_payment_failure_state(event, PaymentError(f"Payment failed after retry: {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 retry count for signing attempts + 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): @@ -232,10 +230,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") @@ -295,47 +294,32 @@ 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) + signed_key = f"payment_signed_{tool_use_id}" + return event.invocation_state.get(signed_key, False) - if retry_count <= 1: - return False + def _mark_successful_signing(self, event: AfterToolCallEvent) -> None: + """Mark that signing succeeded for this tool use. - # 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 + 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. - return False + Args: + event: The after tool call event + """ + tool_use_id = event.tool_use.get("toolUseId", "unknown") + signed_key = f"payment_signed_{tool_use_id}" + event.invocation_state[signed_key] = True @staticmethod def _extract_payment_error_message(body: Optional[Dict[str, Any]]) -> str: diff --git a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py index 9dfb90f9..641c93b7 100644 --- a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py +++ b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py @@ -1310,11 +1310,11 @@ def test_init_agent_passes_none_agent_name_when_not_set(self, mock_payment_manag ) -class TestIsPostPaymentFailure: - """Tests for _is_post_payment_failure method.""" +class TestHasSuccessfulSigning: + """Tests for _has_successful_signing and _mark_successful_signing methods.""" - def test_returns_false_on_first_attempt(self): - """Test that first 402 is not treated as a post-payment failure.""" + def test_returns_false_when_not_signed(self): + """Test that _has_successful_signing returns False when no signing has occurred.""" config = AgentCorePaymentsPluginConfig( payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", user_id="test-user", @@ -1324,33 +1324,14 @@ def test_returns_false_on_first_attempt(self): event, _ = _create_event_with_agent( { "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 1}, - } - ) - - body = {"error": "invalid_exact_evm_insufficient_balance"} - assert plugin._is_post_payment_failure(event, body) is False - - def test_returns_true_on_second_attempt_with_non_payment_error(self): - """Test that a non-'payment required' error on retry is detected as post-payment failure.""" - config = AgentCorePaymentsPluginConfig( - payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", - user_id="test-user", - ) - plugin = AgentCorePaymentsPlugin(config=config) - - event, _ = _create_event_with_agent( - { - "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 2}, + "invocation_state": {}, } ) - body = {"error": "invalid_exact_evm_insufficient_balance"} - assert plugin._is_post_payment_failure(event, body) is True + assert plugin._has_successful_signing(event) is False - def test_returns_false_on_second_attempt_with_payment_required_error(self): - """Test that 'payment required' error on retry is NOT treated as post-payment failure.""" + def test_returns_true_after_marking_signed(self): + """Test that _has_successful_signing returns True after _mark_successful_signing.""" config = AgentCorePaymentsPluginConfig( payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", user_id="test-user", @@ -1360,65 +1341,39 @@ def test_returns_false_on_second_attempt_with_payment_required_error(self): event, _ = _create_event_with_agent( { "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 2}, + "invocation_state": {}, } ) - body = {"error": "Payment required"} - assert plugin._is_post_payment_failure(event, body) is False + plugin._mark_successful_signing(event) + assert plugin._has_successful_signing(event) is True - def test_returns_false_when_body_is_none(self): - """Test that None body is not treated as post-payment failure.""" + def test_signing_state_is_per_tool_use(self): + """Test that signing state is independent per tool_use_id.""" config = AgentCorePaymentsPluginConfig( payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", user_id="test-user", ) plugin = AgentCorePaymentsPlugin(config=config) - event, _ = _create_event_with_agent( + invocation_state = {} + event1, _ = _create_event_with_agent( { - "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 2}, + "tool_use": {"toolUseId": "tool-1"}, + "invocation_state": invocation_state, } ) - - assert plugin._is_post_payment_failure(event, None) is False - - def test_returns_false_when_body_has_empty_error(self): - """Test that empty error string is not treated as post-payment failure.""" - config = AgentCorePaymentsPluginConfig( - payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", - user_id="test-user", - ) - plugin = AgentCorePaymentsPlugin(config=config) - - event, _ = _create_event_with_agent( + event2, _ = _create_event_with_agent( { - "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 2}, + "tool_use": {"toolUseId": "tool-2"}, + "invocation_state": invocation_state, } ) - body = {"error": ""} - assert plugin._is_post_payment_failure(event, body) is False + plugin._mark_successful_signing(event1) - def test_returns_false_when_body_has_no_error_key(self): - """Test that body without error key is not treated as post-payment failure.""" - config = AgentCorePaymentsPluginConfig( - payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", - user_id="test-user", - ) - plugin = AgentCorePaymentsPlugin(config=config) - - event, _ = _create_event_with_agent( - { - "tool_use": {"toolUseId": "tool-123"}, - "invocation_state": {"payment_retry_count_tool-123": 2}, - } - ) - - body = {"statusCode": 402} - assert plugin._is_post_payment_failure(event, body) is False + assert plugin._has_successful_signing(event1) is True + assert plugin._has_successful_signing(event2) is False class TestExtractPaymentErrorMessage: @@ -1453,8 +1408,8 @@ class TestAfterToolCallPostPaymentFailure: """Tests for the post-payment failure path in after_tool_call.""" @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") - def test_post_payment_failure_stores_failure_and_does_not_retry(self, mock_payment_manager_class): - """Test that a 402 with non-payment-required error after retry stores failure without retrying.""" + def test_post_payment_failure_does_not_retry_after_successful_signing(self, mock_payment_manager_class): + """Test that a 402 after successful signing does not retry and leaves result unchanged.""" config = AgentCorePaymentsPluginConfig( payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", user_id="test-user", @@ -1466,20 +1421,20 @@ def test_post_payment_failure_stores_failure_and_does_not_retry(self, mock_payme plugin = AgentCorePaymentsPlugin(config=config) plugin.payment_manager = mock_pm_instance - # Simulate second 402 after a payment retry (retry count already at 1) error_body = { "x402Version": 2, "error": "invalid_exact_evm_insufficient_balance", "resource": {"url": "https://example.com"}, "accepts": [{"scheme": "exact", "network": "eip155:84532", "amount": "1000"}], } + original_result = [ + {"text": f"PAYMENT_REQUIRED: {json.dumps({'statusCode': 402, 'headers': {}, 'body': error_body})}"} + ] event, _ = _create_event_with_agent( { - "result": [ - {"text": f"PAYMENT_REQUIRED: {json.dumps({'statusCode': 402, 'headers': {}, 'body': error_body})}"} - ], + "result": original_result, "tool_use": {"name": "http_request", "toolUseId": "tool-123", "input": {"headers": {}}}, - "invocation_state": {"payment_retry_count_tool-123": 1}, + "invocation_state": {"payment_signed_tool-123": True}, "retry": False, } ) @@ -1488,16 +1443,16 @@ def test_post_payment_failure_stores_failure_and_does_not_retry(self, mock_payme # Should NOT retry assert event.retry is False - # Should store failure state - assert "payment_failure_tool-123" in event.invocation_state - failure = event.invocation_state["payment_failure_tool-123"] - assert "invalid_exact_evm_insufficient_balance" in failure["exceptionMessage"] + # Should NOT store failure state (no interrupt cycle) + assert "payment_failure_tool-123" not in event.invocation_state # Should NOT have called generate_payment_header mock_pm_instance.generate_payment_header.assert_not_called() + # Result should remain unchanged (raw 402 for LLM to reason about) + assert event.result == original_result @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") - def test_first_402_with_payment_required_error_proceeds_normally(self, mock_payment_manager_class): - """Test that the first 402 with 'Payment required' error processes payment normally.""" + def test_first_402_without_prior_signing_proceeds_normally(self, mock_payment_manager_class): + """Test that the first 402 (no prior signing) processes payment normally.""" config = AgentCorePaymentsPluginConfig( payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", user_id="test-user", @@ -1535,6 +1490,72 @@ def test_first_402_with_payment_required_error_proceeds_normally(self, mock_paym # Should retry with payment assert event.retry is True mock_pm_instance.generate_payment_header.assert_called_once() + # Should mark signing as successful + assert event.invocation_state.get("payment_signed_tool-123") is True + + @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") + def test_signing_failure_allows_retry_up_to_max(self, mock_payment_manager_class): + """Test that signing failures are retried up to MAX_PAYMENT_RETRIES.""" + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id="test-user", + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + ) + + mock_pm_instance = MagicMock() + mock_pm_instance.generate_payment_header.side_effect = PaymentError("Signing failed") + plugin = AgentCorePaymentsPlugin(config=config) + plugin.payment_manager = mock_pm_instance + + # Simulate retry count at max - 1 (still below limit) + event, _ = _create_event_with_agent( + { + "result": [{"text": f"PAYMENT_REQUIRED: {json.dumps({'statusCode': 402, 'headers': {}, 'body': {}})}"}], + "tool_use": {"name": "http_request", "toolUseId": "tool-123", "input": {"headers": {}}}, + "invocation_state": {"payment_retry_count_tool-123": 2}, + "retry": False, + } + ) + + plugin.after_tool_call(event) + + # Should attempt signing (even though it fails) + mock_pm_instance.generate_payment_header.assert_called_once() + # Should NOT retry (signing failed) + assert event.retry is False + # Should store failure for interrupt + assert "payment_failure_tool-123" in event.invocation_state + + @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") + def test_signing_retry_limit_stops_processing(self, mock_payment_manager_class): + """Test that after MAX_PAYMENT_RETRIES signing attempts, processing stops.""" + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id="test-user", + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + ) + + mock_pm_instance = MagicMock() + plugin = AgentCorePaymentsPlugin(config=config) + plugin.payment_manager = mock_pm_instance + + # Retry count already at max + event, _ = _create_event_with_agent( + { + "result": [{"text": f"PAYMENT_REQUIRED: {json.dumps({'statusCode': 402, 'headers': {}, 'body': {}})}"}], + "tool_use": {"name": "http_request", "toolUseId": "tool-123", "input": {"headers": {}}}, + "invocation_state": {"payment_retry_count_tool-123": 3}, + "retry": False, + } + ) + + plugin.after_tool_call(event) + + # Should NOT attempt signing + mock_pm_instance.generate_payment_header.assert_not_called() + assert event.retry is False class TestAfterToolCallAutoPaymentDisabled: diff --git a/tests_integ/payments/integrations/strands/test_plugin_integration.py b/tests_integ/payments/integrations/strands/test_plugin_integration.py index 22ffdc5d..0609e0d8 100644 --- a/tests_integ/payments/integrations/strands/test_plugin_integration.py +++ b/tests_integ/payments/integrations/strands/test_plugin_integration.py @@ -710,3 +710,298 @@ def test_plugin_with_agent_name_initializes_payment_manager(self): assert agent is not None logger.info("Plugin with agent_name successfully initialized with real Strands Agent") + + +@pytest.mark.integration +class TestPostPaymentFailureFlow: + """Hook-level integration tests for post-payment failure behavior. + + Exercises the plugin's after_tool_call hook directly to verify the + retry logic end-to-end without requiring an LLM. + """ + + @classmethod + def setup_class(cls): + """Set up test environment.""" + cls.region = os.environ.get("BEDROCK_TEST_REGION", "us-west-2") + cls.user_id = os.environ.get("TEST_USER_ID", f"test-user-{uuid.uuid4().hex[:8]}") + + def _make_402_event(self, tool_use_id, body, invocation_state=None, tool_input=None): + """Create a mock AfterToolCallEvent with a 402 PAYMENT_REQUIRED result.""" + from unittest.mock import MagicMock + + payment_required = { + "statusCode": 402, + "headers": {"content-type": "application/json"}, + "body": body, + } + + event = MagicMock() + event.tool_use = { + "name": "http_request", + "toolUseId": tool_use_id, + "input": tool_input + if tool_input is not None + else {"url": "https://api.example.com/resource", "headers": {}}, + } + event.result = [{"text": f"PAYMENT_REQUIRED: {json.dumps(payment_required)}"}] + event.invocation_state = invocation_state if invocation_state is not None else {} + event.retry = False + event.agent = MagicMock() + event.agent.state.get = MagicMock(return_value=None) + event.agent.state.set = MagicMock() + event.agent.state.delete = MagicMock() + return event + + def test_full_flow_402_sign_retry_402_stops(self): + """Test the full post-payment failure flow through hooks. + + Simulates: + 1. First after_tool_call with 402 → plugin signs and sets retry=True + 2. Second after_tool_call with 402 (server rejection) → plugin stops, no retry + """ + from unittest.mock import MagicMock, patch + + mock_pm = MagicMock() + mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "signed-proof-base64"} + + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id=self.user_id, + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + region=self.region, + ) + + plugin = AgentCorePaymentsPlugin(config=config) + with patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager", return_value=mock_pm): + plugin.init_agent(MagicMock()) + plugin.payment_manager = mock_pm + + # Shared invocation state across retries (persists within a single tool use) + invocation_state = {} + tool_input = {"url": "https://api.example.com/resource", "headers": {}} + + # Step 1: First 402 — payment required + event1 = self._make_402_event( + "tool-abc", + body={"error": "Payment required", "x402Version": 1, "accepts": [{"scheme": "exact"}]}, + invocation_state=invocation_state, + tool_input=tool_input, + ) + + plugin.after_tool_call(event1) + + # Should sign and request retry + assert event1.retry is True + mock_pm.generate_payment_header.assert_called_once() + assert "X-PAYMENT" in tool_input["headers"] + assert invocation_state.get("payment_signed_tool-abc") is True + logger.info("✓ Step 1: First 402 → signed and retry=True") + + # Step 2: Second 402 — server rejected (insufficient balance) + event2 = self._make_402_event( + "tool-abc", + body={"error": "invalid_exact_evm_insufficient_balance", "message": "Insufficient USDC"}, + invocation_state=invocation_state, + tool_input=tool_input, + ) + + plugin.after_tool_call(event2) + + # Should NOT retry — post-payment failure detected + assert event2.retry is False + # generate_payment_header should still only have been called once (from step 1) + assert mock_pm.generate_payment_header.call_count == 1 + # Should NOT store failure state (no interrupt cycle) + assert "payment_failure_tool-abc" not in invocation_state + logger.info("✓ Step 2: Second 402 after signing → stopped, no retry") + + def test_signing_failure_uses_retry_counter(self): + """Test that signing failures increment retry counter and stop at limit. + + Simulates multiple after_tool_call invocations where signing always fails. + Verifies that MAX_PAYMENT_RETRIES is respected. + """ + from unittest.mock import MagicMock, patch + + from bedrock_agentcore.payments.manager import PaymentError + + mock_pm = MagicMock() + mock_pm.generate_payment_header.side_effect = PaymentError("Signing service unavailable") + + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id=self.user_id, + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + region=self.region, + ) + + plugin = AgentCorePaymentsPlugin(config=config) + with patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager", return_value=mock_pm): + plugin.init_agent(MagicMock()) + plugin.payment_manager = mock_pm + + invocation_state = {} + body_402 = {"error": "Payment required", "x402Version": 1, "accepts": [{"scheme": "exact"}]} + + # Attempt signing up to MAX_PAYMENT_RETRIES times + for i in range(plugin.MAX_PAYMENT_RETRIES): + event = self._make_402_event( + "tool-xyz", + body=body_402, + invocation_state=invocation_state, + tool_input={"url": "https://api.example.com", "headers": {}}, + ) + + plugin.after_tool_call(event) + + # Each time signing fails, retry should NOT be set + assert event.retry is False + # Failure should be stored for interrupt + assert "payment_failure_tool-xyz" in invocation_state + # Clean up failure state (simulating interrupt was handled) + del invocation_state["payment_failure_tool-xyz"] + logger.info(" Signing attempt %d/%d failed as expected", i + 1, plugin.MAX_PAYMENT_RETRIES) + + # Signing should have been called exactly MAX_PAYMENT_RETRIES times + assert mock_pm.generate_payment_header.call_count == plugin.MAX_PAYMENT_RETRIES + + # One more attempt should be blocked by retry limit + event_final = self._make_402_event( + "tool-xyz", + body=body_402, + invocation_state=invocation_state, + tool_input={"url": "https://api.example.com", "headers": {}}, + ) + + plugin.after_tool_call(event_final) + + # Should NOT attempt signing — limit reached + assert mock_pm.generate_payment_header.call_count == plugin.MAX_PAYMENT_RETRIES + assert event_final.retry is False + logger.info("✓ Signing retry limit (%d) correctly enforced", plugin.MAX_PAYMENT_RETRIES) + + def test_successful_sign_then_success_response(self): + """Test happy path: 402 → sign → retry succeeds (no second 402). + + After successful signing and retry, if the tool returns a non-402 response, + the plugin should not interfere. + """ + from unittest.mock import MagicMock, patch + + mock_pm = MagicMock() + mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "valid-proof"} + + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id=self.user_id, + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + region=self.region, + ) + + plugin = AgentCorePaymentsPlugin(config=config) + with patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager", return_value=mock_pm): + plugin.init_agent(MagicMock()) + plugin.payment_manager = mock_pm + + invocation_state = {} + tool_input = {"url": "https://api.example.com/resource", "headers": {}} + + # Step 1: 402 → plugin signs and sets retry + event1 = self._make_402_event( + "tool-happy", + body={"error": "Payment required", "x402Version": 1, "accepts": [{"scheme": "exact"}]}, + invocation_state=invocation_state, + tool_input=tool_input, + ) + + plugin.after_tool_call(event1) + + assert event1.retry is True + assert invocation_state.get("payment_signed_tool-happy") is True + assert "X-PAYMENT" in tool_input["headers"] + logger.info("✓ Step 1: Signed and retry set") + + # Step 2: Retry succeeds with 200 — simulate by calling after_tool_call + # with a non-402 result + event2 = MagicMock() + event2.tool_use = {"name": "http_request", "toolUseId": "tool-happy", "input": tool_input} + event2.result = [{"text": json.dumps({"status_code": 200, "body": {"data": "success"}})}] + event2.invocation_state = invocation_state + event2.retry = False + event2.agent = MagicMock() + + plugin.after_tool_call(event2) + + # Should not retry (not a 402) + assert event2.retry is False + # generate_payment_header still only called once + assert mock_pm.generate_payment_header.call_count == 1 + logger.info("✓ Step 2: Non-402 response after retry — plugin did not interfere") + + def test_signing_state_independent_per_tool_use(self): + """Test that signing state is tracked independently per tool_use_id. + + Two different tool uses should have independent signing state: + - tool-A signed successfully → subsequent 402 stops + - tool-B never signed → 402 still attempts signing + """ + from unittest.mock import MagicMock, patch + + mock_pm = MagicMock() + mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "signed"} + + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id=self.user_id, + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + region=self.region, + ) + + plugin = AgentCorePaymentsPlugin(config=config) + with patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager", return_value=mock_pm): + plugin.init_agent(MagicMock()) + plugin.payment_manager = mock_pm + + invocation_state = {} + body_402 = {"error": "Payment required", "x402Version": 1, "accepts": [{"scheme": "exact"}]} + + # tool-A: first 402 → signs successfully + event_a1 = self._make_402_event( + "tool-A", + body=body_402, + invocation_state=invocation_state, + tool_input={"url": "https://a.com", "headers": {}}, + ) + plugin.after_tool_call(event_a1) + assert event_a1.retry is True + assert invocation_state["payment_signed_tool-A"] is True + + # tool-A: second 402 → post-payment failure, stops + event_a2 = self._make_402_event( + "tool-A", + body={"error": "insufficient_balance"}, + invocation_state=invocation_state, + tool_input={"url": "https://a.com", "headers": {}}, + ) + plugin.after_tool_call(event_a2) + assert event_a2.retry is False + + # tool-B: first 402 → should still sign (independent state) + event_b1 = self._make_402_event( + "tool-B", + body=body_402, + invocation_state=invocation_state, + tool_input={"url": "https://b.com", "headers": {}}, + ) + plugin.after_tool_call(event_b1) + assert event_b1.retry is True + assert invocation_state["payment_signed_tool-B"] is True + + # Total signing calls: 2 (one for tool-A, one for tool-B) + assert mock_pm.generate_payment_header.call_count == 2 + logger.info("✓ Signing state correctly independent per tool_use_id") From 16513322767534d523a88ca5f509a89aab13616a Mon Sep 17 00:00:00 2001 From: Raju Ansari Date: Thu, 21 May 2026 11:50:54 -0700 Subject: [PATCH 2/3] fix: address review feedback - Preserve interrupt behavior: post-payment failure stores state and triggers interrupt so downstream consumers are notified - Remove dead code: _extract_payment_error_message (unreachable) - Remove TestExtractPaymentErrorMessage (tested removed method) - Clarify comment on retry counter increment - Add test: signing flag takes precedence over retry limit - Remove emojis from new integration test logger output --- .../payments/integrations/strands/plugin.py | 25 ++---- .../integrations/strands/test_plugin.py | 78 +++++++++++-------- .../strands/test_plugin_integration.py | 10 ++- 3 files changed, 58 insertions(+), 55 deletions(-) diff --git a/src/bedrock_agentcore/payments/integrations/strands/plugin.py b/src/bedrock_agentcore/payments/integrations/strands/plugin.py index 3cf55ab6..485f08a0 100644 --- a/src/bedrock_agentcore/payments/integrations/strands/plugin.py +++ b/src/bedrock_agentcore/payments/integrations/strands/plugin.py @@ -198,12 +198,15 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None: # If we previously signed successfully and still got a 402, the server # rejected the payment for a non-retryable reason (e.g., insufficient balance). - # Stop processing — the existing 402 result is already structured for the LLM. + # 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 successful signing for tool %s — post-payment failure, not retrying", + "Received 402 after successful signing for tool %s — post-payment failure: %s", event.tool_use.get("name", "unknown"), + 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 @@ -211,7 +214,7 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None: logger.warning("Payment signing retry limit reached. Processing skipped.") return - # Increment retry count for signing attempts + # Increment before attempt so limit is enforced even on exception self._increment_payment_retry_count(event) # Validate tool input before processing payment @@ -321,22 +324,6 @@ def _mark_successful_signing(self, event: AfterToolCallEvent) -> None: signed_key = f"payment_signed_{tool_use_id}" event.invocation_state[signed_key] = True - @staticmethod - def _extract_payment_error_message(body: Optional[Dict[str, Any]]) -> str: - """Extract a human-readable error message from a 402 response body. - - Args: - body: The extracted response body (may be None) - - Returns: - Error message string, or "unknown error" if not extractable - """ - if body and isinstance(body, dict): - error = body.get("error") - if isinstance(error, str) and error: - return error - return "unknown error" - def _store_payment_failure_state(self, event: AfterToolCallEvent, exception: Exception) -> None: """Store payment failure information in invocation state for agent to handle. diff --git a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py index 641c93b7..f110bc7b 100644 --- a/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py +++ b/tests/bedrock_agentcore/payments/integrations/strands/test_plugin.py @@ -1376,34 +1376,6 @@ def test_signing_state_is_per_tool_use(self): assert plugin._has_successful_signing(event2) is False -class TestExtractPaymentErrorMessage: - """Tests for _extract_payment_error_message static method.""" - - def test_extracts_error_from_body(self): - """Test extracting error message from body dict.""" - body = {"error": "invalid_exact_evm_insufficient_balance"} - assert AgentCorePaymentsPlugin._extract_payment_error_message(body) == "invalid_exact_evm_insufficient_balance" - - def test_returns_unknown_for_none_body(self): - """Test returns 'unknown error' when body is None.""" - assert AgentCorePaymentsPlugin._extract_payment_error_message(None) == "unknown error" - - def test_returns_unknown_for_empty_error(self): - """Test returns 'unknown error' when error is empty string.""" - body = {"error": ""} - assert AgentCorePaymentsPlugin._extract_payment_error_message(body) == "unknown error" - - def test_returns_unknown_for_missing_error_key(self): - """Test returns 'unknown error' when error key is missing.""" - body = {"statusCode": 402} - assert AgentCorePaymentsPlugin._extract_payment_error_message(body) == "unknown error" - - def test_returns_unknown_for_non_string_error(self): - """Test returns 'unknown error' when error is not a string.""" - body = {"error": 42} - assert AgentCorePaymentsPlugin._extract_payment_error_message(body) == "unknown error" - - class TestAfterToolCallPostPaymentFailure: """Tests for the post-payment failure path in after_tool_call.""" @@ -1443,12 +1415,13 @@ def test_post_payment_failure_does_not_retry_after_successful_signing(self, mock # Should NOT retry assert event.retry is False - # Should NOT store failure state (no interrupt cycle) - assert "payment_failure_tool-123" not in event.invocation_state + # Should store failure state so agent is notified via interrupt + assert "payment_failure_tool-123" in event.invocation_state + failure = event.invocation_state["payment_failure_tool-123"] + assert "Payment rejected after signing" in failure["exceptionMessage"] + assert "invalid_exact_evm_insufficient_balance" in failure["exceptionMessage"] # Should NOT have called generate_payment_header mock_pm_instance.generate_payment_header.assert_not_called() - # Result should remain unchanged (raw 402 for LLM to reason about) - assert event.result == original_result @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") def test_first_402_without_prior_signing_proceeds_normally(self, mock_payment_manager_class): @@ -1558,6 +1531,47 @@ def test_signing_retry_limit_stops_processing(self, mock_payment_manager_class): assert event.retry is False +class TestSigningFlagTakesPrecedenceOverRetryLimit: + """Tests that _has_successful_signing check takes precedence over retry limit.""" + + @patch("bedrock_agentcore.payments.integrations.strands.plugin.PaymentManager") + def test_successful_signing_stops_even_when_retry_limit_not_reached(self, mock_payment_manager_class): + """Test that post-payment failure stops processing regardless of retry counter state. + + If signing succeeded (payment_signed_ flag is True), a subsequent 402 should + stop immediately — even if payment_retry_count is below MAX_PAYMENT_RETRIES. + The signing flag takes precedence over the retry counter. + """ + config = AgentCorePaymentsPluginConfig( + payment_manager_arn="arn:aws:bedrock-agentcore:us-west-2:123456789012:payment-manager/test", + user_id="test-user", + payment_instrument_id="payment-instrument-123", + payment_session_id="payment-session-456", + ) + + mock_pm_instance = MagicMock() + plugin = AgentCorePaymentsPlugin(config=config) + plugin.payment_manager = mock_pm_instance + + # Signing succeeded but retry counter is only at 1 (below MAX_PAYMENT_RETRIES=3) + event, _ = _create_event_with_agent( + { + "result": [{"text": f"PAYMENT_REQUIRED: {json.dumps({'statusCode': 402, 'headers': {}, 'body': {}})}"}], + "tool_use": {"name": "http_request", "toolUseId": "tool-123", "input": {"headers": {}}}, + "invocation_state": {"payment_signed_tool-123": True, "payment_retry_count_tool-123": 1}, + "retry": False, + } + ) + + plugin.after_tool_call(event) + + # Should NOT retry — signing flag takes precedence + assert event.retry is False + mock_pm_instance.generate_payment_header.assert_not_called() + # Should store failure state for interrupt notification + assert "payment_failure_tool-123" in event.invocation_state + + class TestAfterToolCallAutoPaymentDisabled: """Tests for auto_payment=False in after_tool_call.""" diff --git a/tests_integ/payments/integrations/strands/test_plugin_integration.py b/tests_integ/payments/integrations/strands/test_plugin_integration.py index 0609e0d8..638370c5 100644 --- a/tests_integ/payments/integrations/strands/test_plugin_integration.py +++ b/tests_integ/payments/integrations/strands/test_plugin_integration.py @@ -797,7 +797,7 @@ def test_full_flow_402_sign_retry_402_stops(self): mock_pm.generate_payment_header.assert_called_once() assert "X-PAYMENT" in tool_input["headers"] assert invocation_state.get("payment_signed_tool-abc") is True - logger.info("✓ Step 1: First 402 → signed and retry=True") + logger.info("Step 1: First 402 - signed and retry=True") # Step 2: Second 402 — server rejected (insufficient balance) event2 = self._make_402_event( @@ -813,9 +813,11 @@ def test_full_flow_402_sign_retry_402_stops(self): assert event2.retry is False # generate_payment_header should still only have been called once (from step 1) assert mock_pm.generate_payment_header.call_count == 1 - # Should NOT store failure state (no interrupt cycle) - assert "payment_failure_tool-abc" not in invocation_state - logger.info("✓ Step 2: Second 402 after signing → stopped, no retry") + # Should store failure state so agent is notified via interrupt + assert "payment_failure_tool-abc" in invocation_state + failure = invocation_state["payment_failure_tool-abc"] + assert "Payment rejected after signing" in failure["exceptionMessage"] + logger.info("Step 2: Second 402 after signing - stopped, failure stored for interrupt") def test_signing_failure_uses_retry_counter(self): """Test that signing failures increment retry counter and stop at limit. From 743b1a01e8e6b8c32aabe27a0cf380f66c6090f3 Mon Sep 17 00:00:00 2001 From: Aidan Daly Date: Thu, 21 May 2026 20:37:53 +0000 Subject: [PATCH 3/3] fix: address PR review feedback on retry-improvement - plugin.py: drop redundant tool_input re-fetch (already bound at line 177) - plugin.py: document that invocation_state markers are intentionally per-invocation and do not need explicit cleanup - test_plugin_integration.py: hoist unittest.mock and PaymentError imports to module top, per project convention - test_plugin_integration.py: clarify TestPostPaymentFailureFlow docstring to flag it as mock-driven and credential-free --- .../payments/integrations/strands/plugin.py | 6 ++++- .../strands/test_plugin_integration.py | 22 ++++++------------- 2 files changed, 12 insertions(+), 16 deletions(-) diff --git a/src/bedrock_agentcore/payments/integrations/strands/plugin.py b/src/bedrock_agentcore/payments/integrations/strands/plugin.py index 485f08a0..041b53ed 100644 --- a/src/bedrock_agentcore/payments/integrations/strands/plugin.py +++ b/src/bedrock_agentcore/payments/integrations/strands/plugin.py @@ -218,7 +218,6 @@ def after_tool_call(self, event: AfterToolCallEvent) -> None: 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")) @@ -317,6 +316,11 @@ def _mark_successful_signing(self, event: AfterToolCallEvent) -> None: right before setting event.retry. If a subsequent 402 is received, _has_successful_signing will return True indicating the failure is server-side. + 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: event: The after tool call event """ diff --git a/tests_integ/payments/integrations/strands/test_plugin_integration.py b/tests_integ/payments/integrations/strands/test_plugin_integration.py index 638370c5..c49df49d 100644 --- a/tests_integ/payments/integrations/strands/test_plugin_integration.py +++ b/tests_integ/payments/integrations/strands/test_plugin_integration.py @@ -8,6 +8,7 @@ import os import uuid from typing import Any +from unittest.mock import MagicMock, patch import pytest from strands import Agent @@ -16,6 +17,7 @@ from bedrock_agentcore.payments.integrations.config import AgentCorePaymentsPluginConfig from bedrock_agentcore.payments.integrations.strands.plugin import AgentCorePaymentsPlugin from bedrock_agentcore.payments.manager import ( + PaymentError, PaymentInstrumentConfigurationRequired, PaymentManager, PaymentSessionConfigurationRequired, @@ -714,10 +716,12 @@ def test_plugin_with_agent_name_initializes_payment_manager(self): @pytest.mark.integration class TestPostPaymentFailureFlow: - """Hook-level integration tests for post-payment failure behavior. + """Hook-flow scenarios for post-payment failure behavior. - Exercises the plugin's after_tool_call hook directly to verify the - retry logic end-to-end without requiring an LLM. + These are mock-driven tests that exercise the plugin's after_tool_call hook + directly. They fully mock PaymentManager, do not hit AWS, and require no + credentials — they live here alongside the real-AWS integration tests only + because they cover end-to-end hook flow rather than a single helper method. """ @classmethod @@ -728,8 +732,6 @@ def setup_class(cls): def _make_402_event(self, tool_use_id, body, invocation_state=None, tool_input=None): """Create a mock AfterToolCallEvent with a 402 PAYMENT_REQUIRED result.""" - from unittest.mock import MagicMock - payment_required = { "statusCode": 402, "headers": {"content-type": "application/json"}, @@ -760,8 +762,6 @@ def test_full_flow_402_sign_retry_402_stops(self): 1. First after_tool_call with 402 → plugin signs and sets retry=True 2. Second after_tool_call with 402 (server rejection) → plugin stops, no retry """ - from unittest.mock import MagicMock, patch - mock_pm = MagicMock() mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "signed-proof-base64"} @@ -825,10 +825,6 @@ def test_signing_failure_uses_retry_counter(self): Simulates multiple after_tool_call invocations where signing always fails. Verifies that MAX_PAYMENT_RETRIES is respected. """ - from unittest.mock import MagicMock, patch - - from bedrock_agentcore.payments.manager import PaymentError - mock_pm = MagicMock() mock_pm.generate_payment_header.side_effect = PaymentError("Signing service unavailable") @@ -891,8 +887,6 @@ def test_successful_sign_then_success_response(self): After successful signing and retry, if the tool returns a non-402 response, the plugin should not interfere. """ - from unittest.mock import MagicMock, patch - mock_pm = MagicMock() mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "valid-proof"} @@ -951,8 +945,6 @@ def test_signing_state_independent_per_tool_use(self): - tool-A signed successfully → subsequent 402 stops - tool-B never signed → 402 still attempts signing """ - from unittest.mock import MagicMock, patch - mock_pm = MagicMock() mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "signed"}