diff --git a/src/bedrock_agentcore/payments/integrations/config.py b/src/bedrock_agentcore/payments/integrations/config.py index 83bc5c9d..bf17a95b 100644 --- a/src/bedrock_agentcore/payments/integrations/config.py +++ b/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". diff --git a/src/bedrock_agentcore/payments/integrations/langgraph/README.md b/src/bedrock_agentcore/payments/integrations/langgraph/README.md index c0d31b94..161900ba 100644 --- a/src/bedrock_agentcore/payments/integrations/langgraph/README.md +++ b/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): diff --git a/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py b/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py index 16049f4e..10ae637a 100644 --- a/src/bedrock_agentcore/payments/integrations/langgraph/middleware.py +++ b/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) diff --git a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage3.py b/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage3.py index 7bfc0918..3d4928e8 100644 --- a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage3.py +++ b/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 diff --git a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage7.py b/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage7.py index a3ed5113..12dd12a4 100644 --- a/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage7.py +++ b/tests/bedrock_agentcore/payments/integrations/langgraph/test_stage7.py @@ -1,6 +1,7 @@ """Tests for Stage 7: Error Handler Callback.""" import asyncio +import functools import json from unittest.mock import AsyncMock, MagicMock, patch @@ -391,11 +392,11 @@ def test_no_callback_session_expired(self, mock_pm_cls): class TestAsyncCallbackInSyncPath: - """Async callbacks on the sync path fail loudly instead of silently.""" + """Async callbacks on the sync path fail loudly (raise TypeError) instead of silently.""" @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") - def test_async_callback_in_sync_path_falls_through_gracefully(self, mock_pm_cls): - """An async callback on sync .invoke() logs a TypeError and falls through to PROPAGATE.""" + def test_async_callback_in_sync_path_raises(self, mock_pm_cls): + """An async callback on sync .invoke() raises a TypeError instead of a silent PROPAGATE.""" mock_pm_cls.return_value.generate_payment_header.side_effect = PaymentError("fail") async def async_cb(ctx): @@ -408,13 +409,12 @@ async def async_cb(ctx): request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) mock_handler = MagicMock(return_value=ToolMessage(content=_402_content(), tool_call_id="tc-1")) - result = mw.wrap_tool_call(request, mock_handler) - # Falls through to default error message (callback never ran) - assert "PAYMENT ERROR" in result.content + with pytest.raises(TypeError, match="async on_payment_error callback cannot be used with sync"): + mw.wrap_tool_call(request, mock_handler) @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") def test_async_callback_in_sync_path_does_not_mutate_config(self, mock_pm_cls): - """The async callback body never executes, so config is not mutated.""" + """The async callback body never executes (guard raises first), so config is not mutated.""" mock_pm_cls.return_value.generate_payment_header.side_effect = PaymentError("fail") async def async_cb(ctx): @@ -427,9 +427,129 @@ async def async_cb(ctx): request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) mock_handler = MagicMock(return_value=ToolMessage(content=_402_content(), tool_call_id="tc-1")) - mw.wrap_tool_call(request, mock_handler) + with pytest.raises(TypeError): + mw.wrap_tool_call(request, mock_handler) assert config.payment_instrument_id == "instr-1" # unchanged + @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") + def test_async_callable_object_in_sync_path_raises(self, mock_pm_cls): + """A callable object with `async def __call__` is also detected and raises on the sync path.""" + mock_pm_cls.return_value.generate_payment_header.side_effect = PaymentError("fail") + + class AsyncCallback: + async def __call__(self, ctx): + return ErrorResolution.RETRY + + config = _make_config(on_payment_error=AsyncCallback()) + mw = AgentCorePaymentsMiddleware(config) + + request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) + mock_handler = MagicMock(return_value=ToolMessage(content=_402_content(), tool_call_id="tc-1")) + + with pytest.raises(TypeError, match="async on_payment_error callback cannot be used with sync"): + mw.wrap_tool_call(request, mock_handler) + + @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") + def test_async_callable_object_awaited_in_async_path(self, mock_pm_cls): + """A callable object with `async def __call__` is awaited (not leaked) on the async path.""" + mock_pm = mock_pm_cls.return_value + mock_pm.generate_payment_header.side_effect = [ + PaymentError("first fail"), + {"X-PAYMENT": "sig"}, + ] + + class AsyncCallback: + def __init__(self): + self.called = False + + async def __call__(self, ctx): + self.called = True + ctx.config.payment_session_id = "new-sess" + return ErrorResolution.RETRY + + cb = AsyncCallback() + config = _make_config(on_payment_error=cb) + mw = AgentCorePaymentsMiddleware(config) + + request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) + handler = AsyncMock( + side_effect=[ + ToolMessage(content=_402_content(), tool_call_id="tc-1"), + ToolMessage(content=_200_content(), tool_call_id="tc-1"), + ] + ) + + result = asyncio.run(mw.awrap_tool_call(request, handler)) + assert cb.called # the async __call__ actually ran (was awaited) + assert "PAYMENT ERROR" not in result.content + assert json.loads(result.content)["statusCode"] == 200 + + @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") + def test_partial_wrapped_async_callback_in_sync_path_raises(self, mock_pm_cls): + """functools.partial around an async callback is still detected on the sync path.""" + mock_pm_cls.return_value.generate_payment_header.side_effect = PaymentError("fail") + + async def async_cb(ctx, tenant=None): + return ErrorResolution.RETRY + + config = _make_config(on_payment_error=functools.partial(async_cb, tenant="acme")) + mw = AgentCorePaymentsMiddleware(config) + + request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) + mock_handler = MagicMock(return_value=ToolMessage(content=_402_content(), tool_call_id="tc-1")) + + with pytest.raises(TypeError, match="async on_payment_error callback cannot be used with sync"): + mw.wrap_tool_call(request, mock_handler) + + @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") + def test_partial_wrapped_async_object_in_sync_path_raises(self, mock_pm_cls): + """functools.partial around a callable object with async __call__ is also detected. + + This is the case a plain inspect.iscoroutinefunction misses even on Python 3.10. + """ + mock_pm_cls.return_value.generate_payment_header.side_effect = PaymentError("fail") + + class AsyncCallback: + async def __call__(self, ctx, tenant=None): + return ErrorResolution.RETRY + + config = _make_config(on_payment_error=functools.partial(AsyncCallback(), tenant="acme")) + mw = AgentCorePaymentsMiddleware(config) + + request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) + mock_handler = MagicMock(return_value=ToolMessage(content=_402_content(), tool_call_id="tc-1")) + + with pytest.raises(TypeError, match="async on_payment_error callback cannot be used with sync"): + mw.wrap_tool_call(request, mock_handler) + + @patch("bedrock_agentcore.payments.integrations.langgraph.middleware.PaymentManager") + def test_partial_wrapped_async_callback_awaited_in_async_path(self, mock_pm_cls): + """A functools.partial async callback is awaited (with its bound args) on the async path.""" + mock_pm = mock_pm_cls.return_value + mock_pm.generate_payment_header.side_effect = [PaymentError("first fail"), {"X-PAYMENT": "sig"}] + + seen = [] + + async def async_cb(ctx, tenant=None): + seen.append(tenant) + ctx.config.payment_session_id = "new-sess" + return ErrorResolution.RETRY + + config = _make_config(on_payment_error=functools.partial(async_cb, tenant="acme")) + mw = AgentCorePaymentsMiddleware(config) + + request = _make_request(tool_args={"url": "http://x.com", "headers": {}}) + handler = AsyncMock( + side_effect=[ + ToolMessage(content=_402_content(), tool_call_id="tc-1"), + ToolMessage(content=_200_content(), tool_call_id="tc-1"), + ] + ) + + result = asyncio.run(mw.awrap_tool_call(request, handler)) + assert seen == ["acme"] # partial's bound arg applied and awaited + assert json.loads(result.content)["statusCode"] == 200 + # --------------------------------------------------------------------------- # Post-recovery rejection with raw JSON (fallback handler path)