bedrock-agentcore-starter-toolkit 0.1.25__py3-none-any.whl → 0.1.26__py3-none-any.whl

bedrock_agentcore_starter_toolkit/cli/runtime/commands.py

@@ -4,6 +4,7 @@ import json
 import logging
 import os
 from pathlib import Path
+from threading import Thread
 from typing import List, Optional
 
 import typer
@@ -12,6 +13,7 @@ from prompt_toolkit.completion import PathCompleter
 from rich.panel import Panel
 from rich.syntax import Syntax
 
+from ...operations.identity.oauth2_callback_server import start_oauth2_callback_server
 from ...operations.runtime import (
     configure_bedrock_agentcore,
     destroy_bedrock_agentcore,
@@ -575,12 +577,23 @@ def launch(
             _print_success(f"Docker image built: {result.tag}")
             _print_success("Ready to run locally")
             console.print("Starting server at http://localhost:8080")
+            console.print("Starting OAuth2 3LO callback server at http://localhost:8081")
             console.print("[yellow]Press Ctrl+C to stop[/yellow]\n")
 
             if result.runtime is None or result.port is None:
                 _handle_error("Unable to launch locally")
 
             try:
+                oauth2_callback_endpoint = Thread(
+                    target=start_oauth2_callback_server,
+                    args=(
+                        config_path,
+                        agent,
+                    ),
+                    name="OAuth2 3LO Callback Server",
+                    daemon=True,
+                )
+                oauth2_callback_endpoint.start()
                 result.runtime.run_local(result.tag, result.port, result.env_vars)
             except KeyboardInterrupt:
                 console.print("\n[yellow]Stopped[/yellow]")

bedrock_agentcore_starter_toolkit/operations/identity/__init__.py

@@ -0,0 +1,5 @@
+"""Bedrock AgentCore Identity operations."""
+
+from .oauth2_callback_server import WORKLOAD_USER_ID, start_oauth2_callback_server
+
+__all__ = ["start_oauth2_callback_server", "WORKLOAD_USER_ID"]

bedrock_agentcore_starter_toolkit/operations/identity/oauth2_callback_server.py

@@ -0,0 +1,86 @@
+"""Provides a Starlette-based web server that handles OAuth2 3LO callbacks."""
+
+from pathlib import Path
+
+import uvicorn
+from bedrock_agentcore.services.identity import IdentityClient, UserIdIdentifier
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse
+from starlette.routing import Route
+
+from ...cli.common import console
+from ...utils.runtime.config import BedrockAgentCoreAgentSchema, load_config
+
+OAUTH2_CALLBACK_SERVER_PORT = 8081
+OAUTH2_CALLBACK_ENDPOINT = "/oauth2/callback"
+WORKLOAD_USER_ID = "userId"
+
+
+def start_oauth2_callback_server(config_path: Path, agent_name: str, debug: bool = False):
+    """Starts a server to complete the OAuth2 3LO flow with AgentCore Identity."""
+    callback_server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name=agent_name, debug=debug)
+    callback_server.run()
+
+
+class BedrockAgentCoreIdentity3loCallback(Starlette):
+    """Bedrock AgentCore application class that extends Starlette for OAuth2 3LO callback flow."""
+
+    def __init__(self, config_path: Path, agent_name: str, debug: bool = False):
+        """Initialize Bedrock AgentCore Identity callback server."""
+        self.config_path = config_path
+        self.agent_name = agent_name
+        routes = [
+            Route(OAUTH2_CALLBACK_ENDPOINT, self._handle_3lo_callback, methods=["GET"]),
+        ]
+        super().__init__(routes=routes, debug=debug)
+
+    def run(self, **kwargs):
+        """Start the Bedrock AgentCore Identity OAuth2 callback server."""
+        uvicorn_params = {
+            "host": "127.0.0.1",
+            "port": OAUTH2_CALLBACK_SERVER_PORT,
+            "access_log": self.debug,
+            "log_level": "info" if self.debug else "warning",
+        }
+        uvicorn_params.update(kwargs)
+
+        uvicorn.run(self, **uvicorn_params)
+
+    def _handle_3lo_callback(self, request: Request) -> JSONResponse:
+        """Handle OAuth2 3LO callbacks with AgentCore Identity."""
+        session_id = request.query_params.get("session_id")
+        if not session_id:
+            console.print("Missing session_id in OAuth2 3LO callback")
+            return JSONResponse(status_code=400, content={"message": "missing session_id query parameter"})
+
+        project_config = load_config(self.config_path)
+        agent_config: BedrockAgentCoreAgentSchema = project_config.get_agent_config(self.agent_name)
+        oauth2_config = agent_config.oauth_configuration
+
+        user_id = None
+        if oauth2_config:
+            user_id = oauth2_config.get(WORKLOAD_USER_ID)
+
+        if not user_id:
+            console.print(f"Missing {WORKLOAD_USER_ID} in Agent OAuth2 Config")
+            return JSONResponse(status_code=500, content={"message": "Internal Server Error"})
+
+        console.print(f"Handling 3LO callback for workload_user_id={user_id} | session_id={session_id}", soft_wrap=True)
+
+        region = agent_config.aws.region
+        if not region:
+            console.print("AWS Region not configured")
+            return JSONResponse(status_code=500, content={"message": "Internal Server Error"})
+
+        identity_client = IdentityClient(region)
+        identity_client.complete_resource_token_auth(
+            session_uri=session_id, user_identifier=UserIdIdentifier(user_id=user_id)
+        )
+
+        return JSONResponse(status_code=200, content={"message": "OAuth2 3LO flow completed successfully"})
+
+    @classmethod
+    def get_oauth2_callback_endpoint(cls) -> str:
+        """Returns the url for the local OAuth2 callback server."""
+        return f"http://localhost:{OAUTH2_CALLBACK_SERVER_PORT}{OAUTH2_CALLBACK_ENDPOINT}"

bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py

@@ -7,6 +7,7 @@ from typing import Any, Optional
 
 from bedrock_agentcore.services.identity import IdentityClient
 
+from ...operations.identity.oauth2_callback_server import WORKLOAD_USER_ID, BedrockAgentCoreIdentity3loCallback
 from ...services.runtime import BedrockAgentCoreClient, generate_session_id
 from ...utils.runtime.config import load_config, save_config
 from ...utils.runtime.schema import BedrockAgentCoreConfigSchema
@@ -121,9 +122,19 @@ def invoke_bedrock_agentcore(
             workload_name=workload_name, user_token=bearer_token, user_id=user_id
         )["workloadAccessToken"]
 
+        agent_config.oauth_configuration[WORKLOAD_USER_ID] = user_id  # type: ignore : populated by _get_workload_name(...)
+        save_config(project_config, config_path)
+
+        oauth2_callback_url = BedrockAgentCoreIdentity3loCallback.get_oauth2_callback_endpoint()
+        _update_workload_identity_with_oauth2_callback_url(
+            identity_client, workload_name=workload_name, oauth2_callback_url=oauth2_callback_url
+        )
+
         # TODO: store and read port config of local running container
         client = LocalBedrockAgentCoreClient("http://127.0.0.1:8080")
-        response = client.invoke_endpoint(session_id, payload_str, workload_access_token, custom_headers)
+        response = client.invoke_endpoint(
+            session_id, payload_str, workload_access_token, oauth2_callback_url, custom_headers
+        )
 
     else:
         if not agent_arn:
@@ -163,6 +174,24 @@ def invoke_bedrock_agentcore(
     )
 
 
+def _update_workload_identity_with_oauth2_callback_url(
+    identity_client: IdentityClient,
+    workload_name: str,
+    oauth2_callback_url: str,
+) -> None:
+    workload_identity = identity_client.get_workload_identity(name=workload_name)
+    allowed_resource_oauth_2_return_urls = workload_identity.get("allowedResourceOauth2ReturnUrls") or []
+    if oauth2_callback_url in allowed_resource_oauth_2_return_urls:
+        return
+
+    log.info("Updating workload %s with callback url %s", workload_name, oauth2_callback_url)
+
+    identity_client.update_workload_identity(
+        name=workload_name,
+        allowed_resource_oauth_2_return_urls=[*allowed_resource_oauth_2_return_urls, oauth2_callback_url],
+    )
+
+
 def _get_workload_name(
     project_config: BedrockAgentCoreConfigSchema,
     project_config_path: Path,

bedrock_agentcore_starter_toolkit/services/runtime.py

@@ -576,10 +576,11 @@ class LocalBedrockAgentCoreClient:
         session_id: str,
         payload: str,
         workload_access_token: str,
+        oauth2_callback_url: str,
         custom_headers: Optional[dict] = None,
     ):
         """Invoke the endpoint with the given parameters."""
-        from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, SESSION_HEADER
+        from bedrock_agentcore.runtime.models import ACCESS_TOKEN_HEADER, OAUTH2_CALLBACK_URL_HEADER, SESSION_HEADER
 
         url = f"{self.endpoint}/invocations"
 
@@ -587,6 +588,7 @@ class LocalBedrockAgentCoreClient:
             "Content-Type": "application/json",
             ACCESS_TOKEN_HEADER: workload_access_token,
             SESSION_HEADER: session_id,
+            OAUTH2_CALLBACK_URL_HEADER: oauth2_callback_url,
         }
 
         # Merge custom headers if provided

{bedrock_agentcore_starter_toolkit-0.1.25.dist-info → bedrock_agentcore_starter_toolkit-0.1.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bedrock-agentcore-starter-toolkit
-Version: 0.1.25
+Version: 0.1.26
 Summary: A starter toolkit for using Bedrock AgentCore
 Project-URL: Homepage, https://github.com/aws/bedrock-agentcore-starter-toolkit
 Project-URL: Bug Tracker, https://github.com/aws/bedrock-agentcore-starter-toolkit/issues
@@ -22,9 +22,9 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.10
 Requires-Dist: autopep8>=2.3.2
-Requires-Dist: bedrock-agentcore>=0.1.7
-Requires-Dist: boto3>=1.40.35
-Requires-Dist: botocore>=1.40.35
+Requires-Dist: bedrock-agentcore>=1.0.3
+Requires-Dist: boto3>=1.40.51
+Requires-Dist: botocore>=1.40.51
 Requires-Dist: docstring-parser<1.0,>=0.15
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: jinja2>=3.1.6
@@ -32,12 +32,13 @@ Requires-Dist: openapi-spec-validator>=0.7.2
 Requires-Dist: prance>=25.4.8.0
 Requires-Dist: prompt-toolkit>=3.0.51
 Requires-Dist: py-openapi-schema-to-json-schema>=0.0.3
-Requires-Dist: pydantic<3.0.0,>=2.0.0
+Requires-Dist: pydantic<2.41.3,>=2.0.0
 Requires-Dist: pyyaml>=6.0.2
 Requires-Dist: questionary>=2.1.0
 Requires-Dist: requests>=2.25.0
 Requires-Dist: rich>=13.0.0
 Requires-Dist: ruamel-yaml>=0.18.14
+Requires-Dist: starlette>=0.46.2
 Requires-Dist: toml>=0.10.2
 Requires-Dist: typer>=0.16.0
 Requires-Dist: typing-extensions<5.0.0,>=4.13.2

{bedrock_agentcore_starter_toolkit-0.1.25.dist-info → bedrock_agentcore_starter_toolkit-0.1.26.dist-info}/RECORD

@@ -9,7 +9,7 @@ bedrock_agentcore_starter_toolkit/cli/import_agent/__init__.py,sha256=tyM3dXM3Tc
 bedrock_agentcore_starter_toolkit/cli/import_agent/agent_info.py,sha256=CcPXbKNnI9fc7NyFmZBjB3PkB0wAz2Q_kBGoUVcnngA,9736
 bedrock_agentcore_starter_toolkit/cli/import_agent/commands.py,sha256=vpA66fGwKf6cxBelT3Yb_Af5Ow3AMquGLIDe-tBLjmc,22419
 bedrock_agentcore_starter_toolkit/cli/runtime/__init__.py,sha256=0zKPPoYThoDIr3DZhIQJavq5nVTKRsgcE1s9--wx118,60
-bedrock_agentcore_starter_toolkit/cli/runtime/commands.py,sha256=NntVm0PQshWxrB-AJ-pQVo8UAIgxOt5Y3I_DfF4bOlM,55971
+bedrock_agentcore_starter_toolkit/cli/runtime/commands.py,sha256=CLRynotRlQST9gsbbmDg4_Y7JXTZFt_G-sh11OgH7O4,56558
 bedrock_agentcore_starter_toolkit/cli/runtime/configuration_manager.py,sha256=aRHo8tWrx4IBJPfRMQxUe3eZsoOQn1qUYJ_ZI6MAPSo,15832
 bedrock_agentcore_starter_toolkit/notebook/__init__.py,sha256=wH1ZzIVKsKT_P0qX2kIIoyVxeHj8K40odfR1YI3McHw,129
 bedrock_agentcore_starter_toolkit/notebook/runtime/__init__.py,sha256=DoGfB_uGFLANJVE9rSZtTH3ymw76WBWmD9vORvBIdXs,66
@@ -21,6 +21,8 @@ bedrock_agentcore_starter_toolkit/operations/gateway/constants.py,sha256=0_4J6BN
 bedrock_agentcore_starter_toolkit/operations/gateway/create_lambda.py,sha256=mAzXvi5hetzN2iv0ITfx6PDbOgeLOjqt71zzaDpKHnw,2876
 bedrock_agentcore_starter_toolkit/operations/gateway/create_role.py,sha256=Oma9s6K4T9EFgwPAUDungDL9fGa3W0LsHBknIXUmSpI,4474
 bedrock_agentcore_starter_toolkit/operations/gateway/exceptions.py,sha256=WjsrE7lT2708CJniM_ISMzkfNX4IUteGgvOxleD9JZY,272
+bedrock_agentcore_starter_toolkit/operations/identity/__init__.py,sha256=8gMJRmRmHy1_MJs4oJnM3RlkMgfx204EcvyIxRoP3oQ,193
+bedrock_agentcore_starter_toolkit/operations/identity/oauth2_callback_server.py,sha256=2FlCvjxbNOP6d0ciRZ7QIXHmyJIEIN8FSSGBONEYaZc,3669
 bedrock_agentcore_starter_toolkit/operations/memory/README.md,sha256=VHmdCoZImRVsiDvCnKbnt4btO5fOkfeKmoludgSbKwM,31139
 bedrock_agentcore_starter_toolkit/operations/memory/__init__.py,sha256=PAj25QAlnlvjaQIgfIpfxXXVXu5BTOXy-xDMDh6_HIc,59
 bedrock_agentcore_starter_toolkit/operations/memory/constants.py,sha256=0HWpxJZXmnmekIEW5TJjn9KmEqVi_REJzUN1l68ohYQ,3226
@@ -42,14 +44,14 @@ bedrock_agentcore_starter_toolkit/operations/runtime/configure.py,sha256=b226Wg2
 bedrock_agentcore_starter_toolkit/operations/runtime/create_role.py,sha256=1-b_wBvkOXNh-HJsxATwHfXQqj0dpdETds0FNSTY0BE,16116
 bedrock_agentcore_starter_toolkit/operations/runtime/destroy.py,sha256=iyAH1-D734cmiluASGMcp3wwXZ10IlhKi-acsGEQ8Yk,26204
 bedrock_agentcore_starter_toolkit/operations/runtime/exceptions.py,sha256=7iZNVbsz5XyvnpMM4paJb5_LT6DYdEI3nHhM9f3BERE,869
-bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py,sha256=ENrJd6DP30DuXe94_kjrMM3njUUbQMsX7O4rM6hFL5g,7162
+bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py,sha256=t8Uv9E9IVWedspQ3ezx0Sz95LfGVweY7diJ315loGlE,8471
 bedrock_agentcore_starter_toolkit/operations/runtime/launch.py,sha256=oVTLxDdAL2bG-oG4S-KIuUntSODIftit7nGslzjsslg,30100
 bedrock_agentcore_starter_toolkit/operations/runtime/models.py,sha256=TdUS-O1crCGmigKb_N3WaLBeubsuFiUczyZktzcAHjY,4894
 bedrock_agentcore_starter_toolkit/operations/runtime/status.py,sha256=yQR0dOjPeA3jyWg9EcezsTi6OVlp_9tbURe3clQNBaY,5340
 bedrock_agentcore_starter_toolkit/services/__init__.py,sha256=s7QtYYFCCX2ff0gZHUdDxCDI3zhUq0fPsjevZbF9xdA,66
 bedrock_agentcore_starter_toolkit/services/codebuild.py,sha256=CA4OgkMiZPbV6KBFswBLlSuaWgcH1wnTmRnvSEmHxpc,15925
 bedrock_agentcore_starter_toolkit/services/ecr.py,sha256=a9VpxIzYohPvAiAsl9jPp5vZcalCxt9LpqW4S7CdFPo,3991
-bedrock_agentcore_starter_toolkit/services/runtime.py,sha256=lqBB-4uJuxNKbnPHpRskQ7bajiZ2MpjnbqPnKnvTinY,22408
+bedrock_agentcore_starter_toolkit/services/runtime.py,sha256=Ec7l7hj4gRxzSfVJJ2rOwUr3IvuMyDU-u81ueJqF-kI,22531
 bedrock_agentcore_starter_toolkit/services/xray.py,sha256=CRlcfVXpghVy7PvZqLUC4rp49Sw5DQ98rUU61HAluRM,6363
 bedrock_agentcore_starter_toolkit/services/import_agent/__init__.py,sha256=ig-xanZqA3oJdRTucYz9xF9_2yi31ADRVIKFsnIw_l4,68
 bedrock_agentcore_starter_toolkit/services/import_agent/utils.py,sha256=yCnzqv3S8sbQi6ArTz3MeR4ggRcojbAgMbBb7KAf0sA,16077
@@ -73,9 +75,9 @@ bedrock_agentcore_starter_toolkit/utils/runtime/templates/Dockerfile.j2,sha256=7
 bedrock_agentcore_starter_toolkit/utils/runtime/templates/dockerignore.template,sha256=IDKPnAiCdMLjuvLdItzQiybCW--SEvHAN2yr-B5Mxj4,711
 bedrock_agentcore_starter_toolkit/utils/runtime/templates/execution_role_policy.json.j2,sha256=BnVE3ofq2PCrl56VFjilG1RkJEiZ2L3MM2wFok7XTaU,5667
 bedrock_agentcore_starter_toolkit/utils/runtime/templates/execution_role_trust_policy.json.j2,sha256=PPJF6Ofq70W5DUE0NlbmnZjw5RkgepPgkskxEgEG28o,473
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/METADATA,sha256=9CZeflVckSlq6VKgWYaeWJx_1_7aDPk4DB9WfTF7Pdc,9922
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/entry_points.txt,sha256=Tf94DkUf2Tp8P7p8MEXLxre7A7Pp_XNukteiz0wHnk8,77
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/licenses/LICENSE.txt,sha256=nNPOMinitYdtfbhdQgsPgz1UowBznU6QVN3Xs0pSTKU,10768
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/licenses/NOTICE.txt,sha256=rkBsg8DbKqfIoQbveqX9foR4uJPUVAokbkr02pRPilE,8674
-bedrock_agentcore_starter_toolkit-0.1.25.dist-info/RECORD,,
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/METADATA,sha256=RtwR5mxnOFEtLJdPIExD1aoVhudoFvyxwawep7lgPDw,9956
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/entry_points.txt,sha256=Tf94DkUf2Tp8P7p8MEXLxre7A7Pp_XNukteiz0wHnk8,77
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/licenses/LICENSE.txt,sha256=nNPOMinitYdtfbhdQgsPgz1UowBznU6QVN3Xs0pSTKU,10768
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/licenses/NOTICE.txt,sha256=rkBsg8DbKqfIoQbveqX9foR4uJPUVAokbkr02pRPilE,8674
+bedrock_agentcore_starter_toolkit-0.1.26.dist-info/RECORD,,