fix(payments): address langgraph middleware review follow-ups

src/bedrock_agentcore/payments/integrations/config.py

@@ -35,6 +35,12 @@ class AgentCorePaymentsPluginConfig:
             Defaults to 5. Set to 0 to disable interrupt retries entirely.
         custom_handlers: Custom PaymentResponseHandler instances keyed by tool name.
             Takes precedence over the built-in handler registry during resolution.
+            A custom handler's extract_* methods receive the raw ToolMessage.content
+            (a str or a list of content blocks), not the middleware's internal wrapped
+            shape. The built-in handlers (GenericPaymentHandler, HttpRequestPaymentHandler,
+            MCPRequestPaymentHandler) expect a different, normalized shape, so passing one
+            of them directly as a custom handler will not detect 402s; subclass
+            PaymentResponseHandler (or wrap a built-in) and parse the raw content instead.
         auto_session: Whether to auto-create a payment session on first 402 if
             payment_session_id is not set. Default False.
         auto_session_budget: Budget for auto-created sessions (USD). Default "1.00".

src/bedrock_agentcore/payments/integrations/langgraph/README.md

@@ -293,12 +293,15 @@ When no callback is configured, or the callback returns `PROPAGATE`, the LLM rec
 
 Register custom `PaymentResponseHandler` implementations for tools with non-standard output formats. The custom handler is used for **all three phases**: detection, extraction, and injection.
 
+> **Input contract:** a custom handler's `extract_*` methods receive the **raw `ToolMessage.content`** — a `str` or a list of content blocks — exactly as the tool returned it, not the middleware's internal wrapped shape. Parse it yourself. The built-in handlers (`GenericPaymentHandler`, `HttpRequestPaymentHandler`, `MCPRequestPaymentHandler`) expect a different, normalized shape, so passing one of them directly as a custom handler will silently fail to detect 402s — subclass `PaymentResponseHandler` (or wrap a built-in) and parse the raw content.
+
 ```python
 from bedrock_agentcore.payments.integrations.handlers import PaymentResponseHandler
 
 class MyMCPHandler(PaymentResponseHandler):
     def extract_status_code(self, result):
-        # Parse your tool's output format to detect 402
+        # `result` is the raw ToolMessage.content (str or list of blocks).
+        # Parse your tool's output format to detect 402.
         ...
 
     def extract_headers(self, result):

src/bedrock_agentcore/payments/integrations/langgraph/middleware.py

@@ -1,6 +1,7 @@
 """AgentCorePaymentsMiddleware for LangGraph agents."""
 
 import asyncio
+import functools
 import inspect
 import json
 import logging
@@ -209,7 +210,8 @@ def _check_guards(
     def _detect_402(self, request: ToolCallRequest, result: ToolMessage) -> Optional[_DetectionResult]:
         """Run 402 detection. Returns detection context if 402 found, None otherwise."""
         tool_name = request.tool_call["name"]
-        tool_args = request.tool_call.get("args", {})
+        # Store args back on the tool call so an injected payment header reaches the retried handler.
+        tool_args = request.tool_call.setdefault("args", {})
         prepared = self._prepare_for_handler(result.content)
         if prepared is None:
             return None
@@ -307,20 +309,30 @@ def _check_retry_rejection(
         if isinstance(retry_result, Command):
             return retry_result
 
-        retry_prepared = self._prepare_for_handler(retry_result.content)
-        if retry_prepared is None:
-            return None
+        tool_name = request.tool_call["name"]
+        if self.config.custom_handlers and tool_name in self.config.custom_handlers:
+            # Mirror _detect_402: a custom handler owns detection for its tool, so it must
+            # also decide whether the post-payment retry is still a 402. It receives raw content.
+            handler = self.config.custom_handlers[tool_name]
+            retry_status = handler.extract_status_code(retry_result.content)
+            retry_body = handler.extract_body(retry_result.content) if retry_status == 402 else None
+        else:
+            retry_prepared = self._prepare_for_handler(retry_result.content)
+            if retry_prepared is None:
+                return None
+
+            _rh = GenericPaymentHandler()
+            retry_status = _rh.extract_status_code(retry_prepared)
+            if retry_status != 402:
+                fallback = self._fallback_detect_402(retry_result.content)
+                if fallback is not None:
+                    retry_status = 402
+                    _rh = _FallbackHandler(fallback)
+            retry_body = _rh.extract_body(retry_prepared) if retry_status == 402 else None
 
-        _rh = GenericPaymentHandler()
-        retry_status = _rh.extract_status_code(retry_prepared)
-        if retry_status != 402:
-            fallback = self._fallback_detect_402(retry_result.content)
-            if fallback is not None:
-                retry_status = 402
-                _rh = _FallbackHandler(fallback)
         if retry_status == 402:
-            retry_body = _rh.extract_body(retry_prepared) or {}
-            detail = retry_body.get("error", "unknown error") if isinstance(retry_body, dict) else "unknown error"
+            body = retry_body or {}
+            detail = body.get("error", "unknown error") if isinstance(body, dict) else "unknown error"
             return self._error_tool_message(
                 request,
                 PaymentError(f"Payment was signed but rejected {error_context} ({detail})."),
@@ -387,6 +399,23 @@ def _create_auto_session(self) -> None:
         self.config.payment_session_id = session["paymentSessionId"]
         logger.info("auto_session: created session %s", self.config.payment_session_id)
 
+    @staticmethod
+    def _is_async_callback(callback: Any) -> bool:
+        """True if the callback ultimately resolves to an async coroutine function.
+
+        Covers the wrappings a plain ``inspect.iscoroutinefunction`` can miss: callable
+        objects with an ``async def __call__``, ``functools.partial`` (whose async-ness is
+        not reliably visible through the partial across Python versions), and partials of
+        either. (``callable()`` — ruff B004's suggestion — cannot distinguish sync from async.)
+        """
+        target = callback
+        while isinstance(target, functools.partial):
+            target = target.func
+        if inspect.iscoroutinefunction(target):
+            return True
+        call = getattr(target, "__call__", None)  # noqa: B004
+        return inspect.iscoroutinefunction(call)
+
     @staticmethod
     def _error_tool_message(request: ToolCallRequest, exception: Exception) -> ToolMessage:
         """Create a ToolMessage with a deterministic error message for the LLM."""
@@ -505,6 +534,15 @@ def _invoke_error_handler(
         handler: Callable,
     ) -> Optional[Union[ToolMessage, Command]]:
         """Invoke on_payment_error callback and retry if requested."""
+        # Fail loudly: an async callback on the sync path can never be awaited here, so a
+        # silent PROPAGATE would hide a programming error. Raise before the loop (outside the
+        # callback try/except) so the TypeError propagates instead of being swallowed.
+        if self._is_async_callback(self.config.on_payment_error):
+            raise TypeError(
+                "async on_payment_error callback cannot be used with sync .invoke(). "
+                "Use agent.ainvoke() or provide a synchronous callback."
+            )
+
         retry_count = 0
         current_exception = exception
 
@@ -514,11 +552,6 @@ def _invoke_error_handler(
             )
 
             try:
-                if inspect.iscoroutinefunction(self.config.on_payment_error):
-                    raise TypeError(
-                        "async on_payment_error callback cannot be used with sync .invoke(). "
-                        "Use agent.ainvoke() or provide a synchronous callback."
-                    )
                 resolution = self.config.on_payment_error(ctx)
             except Exception as cb_err:
                 logger.error("on_payment_error callback raised: %s", cb_err)
@@ -645,7 +678,7 @@ async def _ainvoke_error_handler(
             )
 
             try:
-                if inspect.iscoroutinefunction(self.config.on_payment_error):
+                if self._is_async_callback(self.config.on_payment_error):
                     resolution = await self.config.on_payment_error(ctx)
                 else:
                     resolution = self.config.on_payment_error(ctx)

tests/bedrock_agentcore/payments/integrations/langgraph/test_stage3.py

@@ -6,6 +6,7 @@
 import pytest
 from langchain.messages import ToolMessage
 
+from bedrock_agentcore.payments.integrations.handlers import PaymentResponseHandler
 from bedrock_agentcore.payments.integrations.langgraph import AgentCorePaymentsConfig
 from bedrock_agentcore.payments.integrations.langgraph.middleware import AgentCorePaymentsMiddleware
 from bedrock_agentcore.payments.manager import (
@@ -15,6 +16,39 @@
 )
 
 
+class _BespokeHandler(PaymentResponseHandler):
+    """Custom handler for a tool with a non-standard 402 shape: {"custom_status": 402, ...}.
+
+    Receives the raw ToolMessage.content (a str) per the custom-handler contract.
+    """
+
+    def extract_status_code(self, result):
+        if isinstance(result, str):
+            try:
+                return json.loads(result).get("custom_status")
+            except (ValueError, TypeError):
+                return None
+        return None
+
+    def extract_headers(self, result):
+        return {}
+
+    def extract_body(self, result):
+        if isinstance(result, str):
+            try:
+                return json.loads(result)
+            except (ValueError, TypeError):
+                return {}
+        return {}
+
+    def validate_tool_input(self, tool_input):
+        return isinstance(tool_input, dict)
+
+    def apply_payment_header(self, tool_input, payment_header):
+        tool_input.setdefault("headers", {}).update(payment_header)
+        return True
+
+
 def _make_config(**overrides):
     defaults = {
         "payment_manager_arn": "arn:aws:bedrock-agentcore:us-east-1:123456789012:payment-manager/pm-1",
@@ -360,3 +394,91 @@ def test_validate_tool_input_fails_returns_error(self, mock_pm_cls):
         result = mw.wrap_tool_call(request, mock_handler)
         assert "PAYMENT ERROR" in result.content
         assert "request format" in result.content
+
+
+# ---------------------------------------------------------------------------
+# Custom-handler retry rejection (custom handler owns post-payment 402 detection)
+# ---------------------------------------------------------------------------
+
+
+class TestCustomHandlerRetryRejection:
+    """A registered custom handler decides whether the post-payment retry is still a 402."""
+
+    @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager")
+    def test_custom_format_rejection_detected_after_payment(self, mock_pm_cls):
+        """Retry still-402 in the custom format is detected and its error detail surfaced."""
+        mock_pm = mock_pm_cls.return_value
+        mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "sig"}
+
+        config = _make_config(custom_handlers={"bespoke_tool": _BespokeHandler()})
+        mw = AgentCorePaymentsMiddleware(config)
+
+        bespoke_402 = json.dumps({"custom_status": 402, "error": "bespoke_reject", "accepts": []})
+        request = _make_request(tool_name="bespoke_tool", tool_args={"url": "http://x.com", "headers": {}})
+        mock_handler = MagicMock(
+            side_effect=[
+                ToolMessage(content=bespoke_402, tool_call_id="tc-1"),  # initial 402
+                ToolMessage(content=bespoke_402, tool_call_id="tc-1"),  # retry still 402
+            ]
+        )
+
+        result = mw.wrap_tool_call(request, mock_handler)
+        assert isinstance(result, ToolMessage)
+        assert "signed but rejected" in result.content
+        assert "bespoke_reject" in result.content  # real detail, not "unknown"
+
+    @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager")
+    def test_custom_format_success_after_payment_passes_through(self, mock_pm_cls):
+        """When the retry succeeds in the custom format, the success result is returned (no false rejection)."""
+        mock_pm = mock_pm_cls.return_value
+        mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "sig"}
+
+        config = _make_config(custom_handlers={"bespoke_tool": _BespokeHandler()})
+        mw = AgentCorePaymentsMiddleware(config)
+
+        bespoke_402 = json.dumps({"custom_status": 402, "error": "pay", "accepts": []})
+        success = ToolMessage(content=json.dumps({"custom_status": 200, "data": "ok"}), tool_call_id="tc-1")
+        request = _make_request(tool_name="bespoke_tool", tool_args={"url": "http://x.com", "headers": {}})
+        mock_handler = MagicMock(side_effect=[ToolMessage(content=bespoke_402, tool_call_id="tc-1"), success])
+
+        result = mw.wrap_tool_call(request, mock_handler)
+        assert result is success
+        assert "PAYMENT ERROR" not in result.content
+
+
+# ---------------------------------------------------------------------------
+# Header injection when the tool call arrives without an "args" key
+# ---------------------------------------------------------------------------
+
+
+class TestHeaderInjectionMissingArgs:
+    """A tool call with no 'args' key still gets the payment header written where the retry sees it."""
+
+    @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager")
+    def test_header_injected_when_args_key_absent(self, mock_pm_cls):
+        mock_pm = mock_pm_cls.return_value
+        mock_pm.generate_payment_header.return_value = {"X-PAYMENT": "sig123"}
+
+        config = _make_config()
+        mw = AgentCorePaymentsMiddleware(config)
+
+        # tool_call dict deliberately has NO "args" key
+        request = MagicMock()
+        request.tool_call = {"name": "http_request", "id": "tc-1"}
+
+        call_count = [0]
+
+        def mock_handler(req):
+            call_count[0] += 1
+            if call_count[0] == 1:
+                return ToolMessage(content=_402_content(), tool_call_id="tc-1")
+            return ToolMessage(content=_200_content(), tool_call_id="tc-1")
+
+        result = mw.wrap_tool_call(request, mock_handler)
+
+        # args was created on request.tool_call and the header injected into it,
+        # so the retried handler would have sent it.
+        assert "args" in request.tool_call
+        assert request.tool_call["args"]["headers"]["X-PAYMENT"] == "sig123"
+        assert call_count[0] == 2
+        assert json.loads(result.content)["statusCode"] == 200