feat(agent): integrate CopilotKit with tools, shared state, and generative UI

- Wire CopilotKitMiddleware and LangGraphAGUIAgent into the LangGraph agent
- Add query_data tool with sample financial CSV for generative chart UI
- Add manage_todos/get_todos tools and AgentState with todos field for shared state
- Add Excalidraw MCP middleware and suggestion
- Add unit test infrastructure for agent tools

Co-Authored-By: Brian Love <brian@liveloveapp.com>
Co-Authored-By: Markus Ecker <markus.ecker@gmail.com>

Author: 3 people

patterns/langgraph-single-agent/Dockerfile

@@ -31,8 +31,9 @@ USER bedrock_agentcore
 
 EXPOSE 8080
 
-# Copy only the agent code
+# Copy agent code
 COPY patterns/langgraph-single-agent/langgraph_agent.py .
+COPY patterns/langgraph-single-agent/tools/ tools/
 COPY patterns/utils/ utils/
 
 # Healthcheck using Python (no extra dependencies needed)

patterns/langgraph-single-agent/langgraph_agent.py

@@ -1,213 +1,198 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-from langgraph.prebuilt import create_react_agent
-from langchain_aws import ChatBedrock
-from langchain_mcp_adapters.client import MultiServerMCPClient
+from __future__ import annotations
+
+import base64
+import json
 import os
-import boto3
+import logging
+
+from ag_ui.core import RunAgentInput, RunErrorEvent, RunFinishedEvent
 from bedrock_agentcore.identity.auth import requires_access_token
 from bedrock_agentcore.runtime import BedrockAgentCoreApp, RequestContext
-import traceback
-
-# Use official LangGraph AWS integration for memory
+from copilotkit import CopilotKitMiddleware, LangGraphAGUIAgent
+from langchain.agents import create_agent
+from langchain_aws import ChatBedrock
+from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph_checkpoint_aws import AgentCoreMemorySaver
 
-from utils.auth import extract_user_id_from_context
 from utils.ssm import get_ssm_parameter
+from tools import query_data, AgentState, todo_tools
 
 app = BedrockAgentCoreApp()
 
-# OAuth2 Credential Provider decorator from AgentCore Identity SDK.
-# Automatically retrieves OAuth2 access tokens from the Token Vault (with caching)
-# or fetches fresh tokens from the configured OAuth2 provider when expired.
-# The provider_name references an OAuth2 Credential Provider registered in AgentCore Identity.
+ACTOR_ID_KEYS = ("actor_id", "actorId", "user_id", "userId", "sub")
+
+SYSTEM_PROMPT = """You are a helpful assistant with access to tools via the Gateway and built-in data tools.
+
+When demonstrating charts, always call the query_data tool first to fetch data from the database before calling any chart tool.
+When managing todos, use manage_todos to update the list and get_todos to read the current list.
+When asked about your tools, list them and explain what they do."""
+
+
+def decode_jwt_sub(authorization_header: str | None) -> str | None:
+    if not authorization_header:
+        return None
+
+    parts = authorization_header.strip().split()
+    if len(parts) != 2 or parts[0].lower() != "bearer":
+        return None
+
+    token_parts = parts[1].split(".")
+    if len(token_parts) < 2:
+        return None
+
+    try:
+        payload = token_parts[1]
+        payload += "=" * ((4 - len(payload) % 4) % 4)
+        decoded = base64.urlsafe_b64decode(payload.encode("utf-8"))
+        sub = json.loads(decoded).get("sub")
+        return sub if isinstance(sub, str) and sub else None
+    except Exception:
+        return None
+
+
+def resolve_actor_id(
+    input_data: RunAgentInput, authorization_header: str | None
+) -> str | None:
+    forwarded_props = (
+        input_data.forwarded_props
+        if isinstance(input_data.forwarded_props, dict)
+        else {}
+    )
+
+    for key in ACTOR_ID_KEYS:
+        value = forwarded_props.get(key)
+        if isinstance(value, str) and value:
+            return value
+
+    return decode_jwt_sub(authorization_header)
+
+
 @requires_access_token(
     provider_name=os.environ["GATEWAY_CREDENTIAL_PROVIDER_NAME"],
     auth_flow="M2M",
-    scopes=[]
+    scopes=[],
 )
 async def _fetch_gateway_token(access_token: str) -> str:
-    """
-    Fetch fresh OAuth2 token for AgentCore Gateway authentication.
-    
-    This is async because it's called with 'await' in create_gateway_mcp_client().
-    The @requires_access_token decorator handles token retrieval and refresh:
-    1. Token Retrieval: Calls GetResourceOauth2Token API to fetch token from Token Vault
-    2. Automatic Refresh: Uses refresh tokens to renew expired access tokens
-    3. Error Orchestration: Handles missing tokens and OAuth flow management
-    
-    For M2M (Machine-to-Machine) flows, the decorator uses Client Credentials grant type.
-    The provider_name must match the Name field in the CDK OAuth2CredentialProvider resource.
-    """
     return access_token
 
 
 async def create_gateway_mcp_client() -> MultiServerMCPClient:
-    """
-    Create an MCP client connected to the AgentCore Gateway with OAuth2 authentication.
-
-    MCP (Model Context Protocol) is how agents communicate with tool providers.
-    This creates a client that can talk to the AgentCore Gateway using OAuth2
-    authentication. The Gateway then provides access to Lambda-based tools.
-    
-    This implementation avoids the "closure trap" by calling _fetch_gateway_token()
-    on every invocation of create_gateway_mcp_client(). Since this function is called
-    per-request in agent_stream(), it ensures fresh tokens for each request.
-    """
-    stack_name = os.environ.get('STACK_NAME')
+    stack_name = os.environ.get("STACK_NAME")
     if not stack_name:
         raise ValueError("STACK_NAME environment variable is required")
-    
-    # Validate stack name format to prevent injection
-    if not stack_name.replace('-', '').replace('_', '').isalnum():
+
+    if not stack_name.replace("-", "").replace("_", "").isalnum():
         raise ValueError("Invalid STACK_NAME format")
-    
-    print(f"[AGENT] Creating Gateway MCP client for stack: {stack_name}")
-    
-    # Fetch Gateway URL from SSM
-    gateway_url = get_ssm_parameter(f'/{stack_name}/gateway_url')
-    print(f"[AGENT] Gateway URL from SSM: {gateway_url}")
-    
-    # Fetch fresh token on every call to avoid closure trap
+
+    gateway_url = get_ssm_parameter(f"/{stack_name}/gateway_url")
     fresh_token = await _fetch_gateway_token()
-    
-    # Create MCP client with Bearer token authentication
-    gateway_client = MultiServerMCPClient({
-        "gateway": {
-            "transport": "streamable_http",
-            "url": gateway_url,
-            "headers": {
-                "Authorization": f"Bearer {fresh_token}"
+
+    return MultiServerMCPClient(
+        {
+            "gateway": {
+                "transport": "streamable_http",
+                "url": gateway_url,
+                "headers": {
+                    "Authorization": f"Bearer {fresh_token}",
+                },
             }
         }
-    })
-    
-    print("[AGENT] Gateway MCP client created successfully")
-    return gateway_client
-
-
-async def create_langgraph_agent(user_id: str, session_id: str, tools: list):
-    """
-    Create a LangGraph agent with AgentCore Gateway MCP tools and memory integration.
-    
-    This function sets up a LangGraph StateGraph that can access tools through
-    the AgentCore Gateway and maintains conversation memory.
-    """
-    system_prompt = """You are a helpful assistant with access to tools via the Gateway.
-    When asked about your tools, list them and explain what they do."""
-
-    # Create Bedrock model
-    bedrock_model = ChatBedrock(
+    )
+
+
+def _build_model(streaming: bool) -> ChatBedrock:
+    return ChatBedrock(
         model_id="us.anthropic.claude-sonnet-4-5-20250929-v1:0",
         temperature=0.1,
-        streaming=True
+        max_tokens=16384,
+        streaming=streaming,
+        beta_use_converse_api=True,
     )
 
-    # Get and validate Memory ID
+
+def _build_checkpointer() -> AgentCoreMemorySaver:
     memory_id = os.environ.get("MEMORY_ID")
     if not memory_id:
         raise ValueError("MEMORY_ID environment variable is required")
-    
-    # Configure AgentCore Memory using official LangGraph AWS integration
-    checkpointer = AgentCoreMemorySaver(
+
+    return AgentCoreMemorySaver(
         memory_id=memory_id,
-        region_name=os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
+        region_name=os.environ.get("AWS_DEFAULT_REGION", "us-east-1"),
     )
 
+
+async def create_langgraph_agent(tools: list):
     try:
-        print("[AGENT] Creating LangGraph agent with Gateway tools...")
-        
-        graph = create_react_agent(
-            model=bedrock_model,
-            tools=tools,
-            checkpointer=checkpointer,
-            prompt=system_prompt
+        return create_agent(
+            model=_build_model(streaming=True),
+            tools=[*tools, query_data, *todo_tools],  # MCP tools + data + todo tools
+            checkpointer=_build_checkpointer(),
+            middleware=[CopilotKitMiddleware()],
+            system_prompt=SYSTEM_PROMPT,
+            state_schema=AgentState,  # extends BaseAgentState with todos: list[Todo]
         )
-        
-        print("[AGENT] Agent created successfully with Gateway tools")
-        return graph
-        
-    except Exception as e:
-        print(f"[AGENT ERROR] Error creating LangGraph agent: {e}")
-        print(f"[AGENT ERROR] Exception type: {type(e).__name__}")
-        print(f"[AGENT ERROR] Traceback:")
-        traceback.print_exc()
+    except Exception:
+        logging.exception("Error creating LangGraph agent")
         raise
 
 
+async def create_runtime_graph():
+    mcp_client = await create_gateway_mcp_client()
+    tools = await mcp_client.get_tools()
+    return await create_langgraph_agent(tools)
+
+
+class ActorAwareLangGraphAgent(LangGraphAGUIAgent):
+    async def run(self, input: RunAgentInput):
+        actor_id = (
+            self.config.get("configurable", {}).get("actor_id") if self.config else None
+        )
+        if not actor_id:
+            raise ValueError(
+                "Missing actor identity. Provide forwardedProps.actor_id/user_id "
+                "or include sub claim in the bearer token."
+            )
+
+        self.graph = await create_runtime_graph()
+        self.config = {"configurable": {"actor_id": actor_id}}
+        async for event in super().run(input):
+            yield event
+
+
 @app.entrypoint
-async def agent_stream(payload, context: RequestContext):
-    """
-    Main entrypoint for the LangGraph agent using streaming with Gateway integration.
-    
-    This is the function that AgentCore Runtime calls when the agent receives a request.
-    It extracts the user's query from the payload, securely obtains the user ID from
-    the validated JWT token in the request context, creates a LangGraph agent with Gateway
-    tools and memory, and streams the response back. This function handles the complete
-    request lifecycle with token-level streaming. The user ID is extracted from the 
-    JWT token (via RequestContext).
-    """
-    user_query = payload.get("prompt")
-    session_id = payload.get("runtimeSessionId")
-    
-    if not all([user_query, session_id]):
-        yield {
-            "status": "error",
-            "error": "Missing required fields: prompt or runtimeSessionId"
-        }
-        return
-    
+async def invocations(payload: dict, context: RequestContext):
+    input_data = RunAgentInput.model_validate(payload)
+    authorization_header = None
+    if context.request_headers:
+        authorization_header = context.request_headers.get("Authorization")
+
+    actor_id = resolve_actor_id(input_data, authorization_header)
+    if not actor_id:
+        raise ValueError(
+            "Missing actor identity. Provide forwardedProps.actor_id/user_id "
+            "or include sub claim in the bearer token."
+        )
+
+    request_agent = ActorAwareLangGraphAgent(
+        name="LangGraphSingleAgent",
+        description="LangGraph single agent exposed via AG-UI",
+        graph=None,
+        config={"configurable": {"actor_id": actor_id}},
+    )
+
     try:
-        # Extract user ID securely from the validated JWT token
-        # instead of trusting the payload body (which could be manipulated)
-        user_id = extract_user_id_from_context(context)
-
-        print(f"[STREAM] Starting streaming invocation for user: {user_id}, session: {session_id}")
-        print(f"[STREAM] Query: {user_query}")
-        
-        # Get OAuth2 access token and create Gateway MCP client
-        # The @requires_access_token decorator handles token fetching automatically
-        print("[STREAM] Creating Gateway MCP client (decorator handles OAuth2)...")
-        mcp_client = await create_gateway_mcp_client()
-        print("[STREAM] Gateway MCP client created successfully")
-        
-        print("[STREAM] Loading Gateway tools...")
-        tools = await mcp_client.get_tools()
-        print(f"[STREAM] Loaded {len(tools)} tools from Gateway")
-        
-        # Create agent with the loaded tools
-        graph = await create_langgraph_agent(user_id, session_id, tools)
-        
-        # Configure streaming with actor_id and thread_id for memory
-        config = {
-            "configurable": {
-                "thread_id": session_id,
-                "actor_id": user_id
-            }
-        }
-        
-        # Stream messages using LangGraph's astream with stream_mode="messages"
-        async for event in graph.astream(
-            {"messages": [("user", user_query)]},
-            config=config,
-            stream_mode="messages"
-        ):
-            # event is a tuple: (message_chunk, metadata)
-            message_chunk, metadata = event
-            yield message_chunk.model_dump()
-        
-        print("[STREAM] Streaming completed successfully")
-            
-    except Exception as e:
-        error_msg = str(e) if str(e) else f"{type(e).__name__}: {repr(e)}"
-        print(f"[STREAM ERROR] Error in agent_stream: {error_msg}")
-        print(f"[STREAM ERROR] Exception type: {type(e).__name__}")
-        traceback.print_exc()
-        yield {
-            "status": "error",
-            "error": error_msg
-        }
+        async for event in request_agent.run(input_data):
+            if event is not None:
+                yield event.model_dump(mode="json", by_alias=True, exclude_none=True)
+    except Exception as exc:
+        logging.exception("Agent run failed")
+        yield RunErrorEvent(
+            message=str(exc) or type(exc).__name__,
+            code=type(exc).__name__,
+        ).model_dump(mode="json", by_alias=True, exclude_none=True)
 
 
 if __name__ == "__main__":

patterns/langgraph-single-agent/requirements.txt

@@ -1,8 +1,14 @@
 # LangGraph agent dependencies with pinned versions
+fastapi==0.115.12
+uvicorn==0.34.2
+ag-ui-protocol>=0.1.10
+ag-ui-langgraph==0.0.28
+copilotkit==0.1.83
+partialjson>=0.0.8
 langgraph==1.0.10rc1
+langchain>=0.3.0
 langchain-aws==1.0.0
 langchain-mcp-adapters==0.1.13
-langgraph-checkpoint-aws==1.0.1
+langgraph-checkpoint-aws==1.0.5
 mcp==1.23.1
 bedrock-agentcore==1.0.6
-PyJWT[crypto]>=2.10.1