[luyuzhe111]

Description

Fix sagemaker_rft_handler failing when an AgentCore entrypoint accepts (payload, context). The SageMaker integration guide requires applying this decorator to AgentCore entrypoints. The official AgentCore Python SDK supports the context parameter and covers it in both sync and async tests.

For example:

from bedrock_agentcore.runtime import BedrockAgentCoreApp
from sagemaker.train.rft import sagemaker_rft_handler

app = BedrockAgentCoreApp()

@app.entrypoint
@sagemaker_rft_handler
async def invoke(payload, context):
    return {"reward": 1.0, "session_id": context.session_id}

AgentCore uses inspect.signature() to decide whether to pass the request context. Because the RFT decorator uses functools.wraps, inspection exposes the original two-argument signature, but the wrapper accepts only payload. Invoking the example therefore returns HTTP 500 with TypeError: invoke() takes 1 positional argument but 2 were given, before the handler or rollout feedback logic runs.

Forward additional positional and keyword arguments through both RFT wrappers while retaining wraps. This lets AgentCore supply the context and preserves existing single-argument handlers. Add unit coverage for sync/async handlers, positional and keyword context forwarding, feedback reporting, and rollout context cleanup. Update the decorator's argument documentation and correct its name in the existing example.

Validation

• pytest sagemaker-train/tests/unit/train/test_rft_decorators.py: 8 passed.
• Local compatibility checks with the official BedrockAgentCoreApp, with feedback HTTP requests mocked: sync and async context handlers fail before the fix and succeed after it.
• Black and flake8 checks on the changed Python files: passed.
• Focused Sphinx build of the modified docstring with warnings treated as errors: passed.

[luyuzhe111]

cc @992X @SifeiLi for review. issue identified by @qingliz in awslabs/agentcore-rl-toolkit#209