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

Author: aidandaly24

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
         """

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"}