Implement 3LO Server on localhost:8081 to handle generating OAuth2 tokens (#282)

SkyWalker-CMD · Sherif Riad · web-flow · commit f2d33a5c5d6a · 2025-10-16T19:44:09.000-04:00
* Implement 3LO Server on localhost:8081 to handle generating OAuth2 tokens

* Update uv.lock with the latest AgentCore SDK

* Add more unit test coverage for the new OAuth2 callback server

* Add more unit test coverage to the OAuth2 3LO Server

* Bump BedrockAgent Core SDK client to version 1.0.3

* Add more unit test coverage for the OAuth2 callback server

* Add more unit test coverage for the local OAuth2 callback server

* Add more unit test coverage for local OAuth2 callback server

---------

Co-authored-by: Sherif Riad &lt;sherifri@amazon.com&gt;
diff --git a/pyproject.toml b/pyproject.toml
@@ -26,9 +26,9 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
-    "boto3>=1.40.35",
-    "botocore>=1.40.35",
-    "bedrock-agentcore>=0.1.7",
+    "boto3>=1.40.51",
+    "botocore>=1.40.51",
+    "bedrock-agentcore>=1.0.3",
     "docstring_parser>=0.15,<1.0",
     "httpx>=0.28.1",
     "jinja2>=3.1.6",
diff --git a/src/bedrock_agentcore_starter_toolkit/cli/runtime/commands.py b/src/bedrock_agentcore_starter_toolkit/cli/runtime/commands.py
@@ -4,6 +4,7 @@
 import logging
 import os
 from pathlib import Path
+from threading import Thread
 from typing import List, Optional
 
 import typer
@@ -12,6 +13,7 @@
 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]")
diff --git a/src/bedrock_agentcore_starter_toolkit/operations/identity/__init__.py b/src/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"]
diff --git a/src/bedrock_agentcore_starter_toolkit/operations/identity/oauth2_callback_server.py b/src/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}"
diff --git a/src/bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py b/src/bedrock_agentcore_starter_toolkit/operations/runtime/invoke.py
@@ -7,6 +7,7 @@
 
 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,
diff --git a/src/bedrock_agentcore_starter_toolkit/services/runtime.py b/src/bedrock_agentcore_starter_toolkit/services/runtime.py
@@ -576,17 +576,19 @@ def invoke_endpoint(
         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"
 
         headers = {
             "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
diff --git a/tests/operations/identity/test_oauth2_callback_server.py b/tests/operations/identity/test_oauth2_callback_server.py
@@ -0,0 +1,104 @@
+from unittest.mock import Mock, patch
+
+from bedrock_agentcore.services.identity import UserIdIdentifier
+from starlette.testclient import TestClient
+
+from bedrock_agentcore_starter_toolkit.operations.identity.oauth2_callback_server import (
+    OAUTH2_CALLBACK_ENDPOINT,
+    WORKLOAD_USER_ID,
+    BedrockAgentCoreIdentity3loCallback,
+)
+from bedrock_agentcore_starter_toolkit.utils.runtime.config import save_config
+from bedrock_agentcore_starter_toolkit.utils.runtime.schema import (
+    AWSConfig,
+    BedrockAgentCoreAgentSchema,
+    BedrockAgentCoreConfigSchema,
+    NetworkConfiguration,
+    ObservabilityConfig,
+)
+
+
+def create_test_config(tmp_path, *, agent_name="test-agent", user_id="test-user-id", region="us-west-2"):
+    config_path = tmp_path / ".bedrock_agentcore.yaml"
+
+    agent_config = BedrockAgentCoreAgentSchema(
+        name=agent_name,
+        entrypoint="test_agent.py",
+        container_runtime="docker",
+        aws=AWSConfig(
+            region=region,
+            account="123456789012",
+            execution_role=None,
+            execution_role_auto_create=True,
+            ecr_repository=None,
+            ecr_auto_create=True,
+            network_configuration=NetworkConfiguration(),
+            observability=ObservabilityConfig(),
+        ),
+        oauth_configuration={WORKLOAD_USER_ID: user_id} if user_id else {},
+    )
+
+    project_config = BedrockAgentCoreConfigSchema(default_agent=agent_name, agents={agent_name: agent_config})
+    save_config(project_config, config_path)
+
+    return config_path
+
+
+class TestBedrockAgentCoreIdentity3loCallback:
+    def test_init(self, tmp_path):
+        config_path = create_test_config(tmp_path)
+        server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name="test-agent")
+
+        assert server.config_path == config_path
+        assert server.agent_name == "test-agent"
+        assert len(server.routes) == 1
+        assert server.routes[0].path == OAUTH2_CALLBACK_ENDPOINT
+
+    def test_get_callback_endpoint(self):
+        endpoint = BedrockAgentCoreIdentity3loCallback.get_oauth2_callback_endpoint()
+        assert endpoint == "http://localhost:8081/oauth2/callback"
+
+    def test_handle_3lo_callback_missing_session_id(self, tmp_path):
+        config_path = create_test_config(tmp_path)
+        server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name="test-agent")
+        client = TestClient(server)
+        response = client.get(OAUTH2_CALLBACK_ENDPOINT)
+
+        assert response.status_code == 400
+        assert response.json().get("message") == "missing session_id query parameter"
+
+    @patch("bedrock_agentcore_starter_toolkit.operations.identity.oauth2_callback_server.IdentityClient")
+    def test_handle_3lo_callback_success(self, mock_identity_client, tmp_path):
+        config_path = create_test_config(tmp_path)
+        server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name="test-agent")
+
+        mock_client_instance = Mock()
+        mock_identity_client.return_value = mock_client_instance
+
+        client = TestClient(server)
+        response = client.get(f"{OAUTH2_CALLBACK_ENDPOINT}?session_id=test-session-123")
+
+        assert response.status_code == 200
+        assert response.json().get("message") == "OAuth2 3LO flow completed successfully"
+        mock_identity_client.assert_called_once_with("us-west-2")
+        mock_client_instance.complete_resource_token_auth.assert_called_once_with(
+            session_uri="test-session-123", user_identifier=UserIdIdentifier(user_id="test-user-id")
+        )
+
+    def test_handle_3lo_callback_missing_user_id(self, tmp_path):
+        config_path = create_test_config(tmp_path, user_id="")
+        server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name="test-agent")
+        client = TestClient(server)
+        response = client.get(f"{OAUTH2_CALLBACK_ENDPOINT}?session_id=test-session-123")
+
+        assert response.status_code == 500
+        assert response.json().get("message") == "Internal Server Error"
+
+    def test_handle_3lo_callback_missing_region(self, tmp_path):
+        config_path = create_test_config(tmp_path, region="")
+        server = BedrockAgentCoreIdentity3loCallback(config_path=config_path, agent_name="test-agent")
+        client = TestClient(server)
+        response = client.get(f"{OAUTH2_CALLBACK_ENDPOINT}?session_id=test-session-123")
+
+        assert response.status_code == 500
+        assert response.json().get("message") == "Internal Server Error"
diff --git a/tests/operations/runtime/test_invoke.py b/tests/operations/runtime/test_invoke.py
diff --git a/tests/services/test_runtime.py b/tests/services/test_runtime.py
diff --git a/uv.lock b/uv.lock
