I've been building AI agents for a while now, and streaming responses to a UI has always been the painful part. In previous projects I tried API Gateway streaming, Lambda response streaming, and even AppSync Events via an agent tool call to notify the UI. I also looked at adding my own WebSocket API through API Gateway, which requires managing $connect
, $disconnect
, and $default
routes, storing connection IDs in DynamoDB, and posting messages back through @connections
. All of these approaches felt like too much ceremony for what should be simple.
That's when I found that AgentCore has built-in WebSocket support. The browser can just connect directly to the agent. No middleman.
I built WearCast to demonstrate this functionality. This is an AI agent that helps you pick out what to where based on the weather. I've found this very helpfu when packing for a trip. The code for this application is public and you can find the full implementation on this GitHub repo.
We usually build app applications as we do REST APIs
User → REST API → Lambda → AI Service → Lambda → REST API → User
Every message makes a round trip through multiple intermediaries. The response waits until the entire generation is complete, and then it all comes back at once. For a chat interface, this feels sluggish. Users stare at a spinner while the model generates hundreds of tokens they could already be reading.
Even if you add Server-Sent Events or long-polling, you're still stitching together a real-time experience on top of infrastructure that seemed like overkill. I wanted something better.
Here's what I ended up with:
┌─────────────┐ ┌──────────────────┐
│ React UI │─── JWT ─▶│ API Gateway + │──▶ Lambda (presigned URL)
│ (Cognito) │ │ Cognito Auth │
└──────┬──────┘ └──────────────────┘
│
│ WebSocket (SigV4 presigned URL)
│ ← No middleman! Direct connection →
▼
┌──────────────────────────────────┐ ┌────────────────────┐
│ AgentCore Runtime │────▶│ AgentCore Memory │
│ (Strands Agent + Bedrock LLM) │ │ (Persistence) │
└──────────────────────────────────┘ └────────────────────┘
The browser connects directly to the AI agent over a WebSocket without the need for any Lambda functions to proxy the tokens. The agent streams tokens directly to the user's browser as they're generated.
The only additional infrastructure is during the initial handshake, where we exchange a JWT for a presigned WebSocket URL. After that, it's a direct, bidirectional pipe.
Let's go over the three steps to get this working.
The user signs in through Amazon Cognito and receives a JWT token. Nothing unusual here.
// Frontend authenticates and gets an access token
const accessToken = await authService.getAccessToken();
Here's where it gets interesting. The frontend makes a single REST call to our backend:
const presignedData = await apiService.getPresignedWebSocketUrl(sessionId, accessToken);
Behind the scenes, a Lambda function:
// Lambda: Generate presigned WebSocket URL
const wsHost = `bedrock-agentcore.${region}.amazonaws.com`;
const wsPath = `/runtimes/${runtimeArn}/ws`;
const signer = new SignatureV4({
service: 'bedrock-agentcore',
region,
credentials,
sha256: Sha256
});
const signedRequest = await signer.presign(request, { expiresIn: 300 });
const presignedWsUrl = formatSignedUrl(signedRequest).replace('https://', 'wss://');
The presigned URL is valid for 5 minutes, long enough to establish the connection, short enough to limit exposure. After the URL expires, it can reconnect by getting a new presigned URL, similar to how we handle expiring authentication tokens.
The browser opens a WebSocket connection using the presigned URL. No custom headers needed, all authentication is embedded in the URL's query parameters via SigV4.
this.ws = new WebSocket(presignedData.wsUrl);
Once connected, the browser sends messages directly to the agent and receives streaming responses in real-time:
// Send a message
ws.send(JSON.stringify({
request: "What should I wear in Chicago today?",
session_id: sessionId,
user_id: userId
}));
// Receive streaming tokens
ws.onmessage = (event) => {
const data = JSON.parse(event.data);
if (data.event?.data) {
// Append token to the UI immediately
appendToStream(data.event.data);
}
};
That's it. The browser is now talking directly to the AI agent. Every token streams in as it's generated.
Let me go over what makes this better than the traditional setup.
Every token arrives at the browser the moment the model generates it. There's no buffering layer, no Lambda invocation overhead per chunk, no API Gateway.
As I mentioned earlier, I've done the traditional API Gateway WebSocket approach and the AppSync Events approach. Both work, but they have a lot of moving parts where they are not needed.
With this approach all you need is:
With this approach you remove the need for any extra components that might generate cost. You only pay for the AgentCore Runtime and one Lambda invocation per session establishment.
The WebSocket stays open for multi-turn conversations. The agent maintains state across messages within the same connection, without the need to reload context on every request.
A direct WebSocket connection is great, but what happens when the user steps away or opens the conversation on a different device? Without memory, the agent starts fresh every time.
WearCast uses AgentCore Memory, a module provided by AgentCore to keep a history of the conversation to give the agent better context awareness
from bedrock_agentcore.memory.integrations.strands.config import AgentCoreMemoryConfig
from bedrock_agentcore.memory.integrations.strands.session_manager import AgentCoreMemorySessionManager
def create_session_manager(runtime_session_id, user_id):
config = AgentCoreMemoryConfig(
memory_id=AGENTCORE_MEMORY_ID,
session_id=runtime_session_id,
actor_id=user_id
)
return AgentCoreMemorySessionManager(
agentcore_memory_config=config,
region_name=AWS_REGION
)
In this case we are using Strands session manager functionality which takes care of automatically previous messages from memory.
New messages are added to the context as they are happening. With this the user can close their laptop, come back hours later, and pick up right where they left off.
The memory is scoped by session ID and actor ID (user), so each user's conversations are isolated and private.
The memory resource is declared right alongside the runtime in the SAM template:
AgentCoreShortTermMemory:
Type: AWS::BedrockAgentCore::Memory
Properties:
Name: WearCast
Description: Short-term memory for agent conversation persistence
MemoryExecutionRoleArn: !GetAtt AgentCoreRole.Arn
EventExpiryDuration: 30 # days
The memory ID is passed to the agent as an environment variable for the session manager to use.
An agent without tools is just a chatbot. Tools turn it into something that can actually do things.
WearCast includes a get_weather
tool that fetches real forecast data from Open-Meteo (no API key required):
@tool
def get_weather(city: str, date: str = "today") -> dict:
"""Get weather conditions for a city, current or up to 16 days ahead.
Args:
city: City name (e.g. "Indianapolis", "Chicago")
date: "today" for current, or YYYY-MM-DD for forecast
"""
geo_url = f"https://geocoding-api.open-meteo.com/v1/search?name={city}&count=1"
forecast_url = f"https://api.open-meteo.com/v1/forecast?latitude={lat}&longitude={lon}&..."
With Strands, tools are just decorated Python functions. The @tool
decorator handles:
On the frontend, tool usage is communicated through the same WebSocket stream:
if (event.current_tool_use?.name) {
setCurrentTool(event.current_tool_use.name);
// Show "Using get_weather..." indicator
}
The user sees the agent "thinking," then using a tool, then formulating its response, all streaming in real-time. It feels like watching someone work, not waiting for an answer.
You might be asking yourself, isn't it dangerous to let browsers connect directly to your AI agent? Not when you layer the security correctly:
// User identity travels with the signed URL
queryParams['X-Amzn-Bedrock-AgentCore-Runtime-Custom-User-Id'] = userId;
// The agent receives it as a header
headers = context.request_headers
user_id = headers.get("x-amzn-bedrock-agentcore-runtime-custom-user-id")
The browser never sees AWS credentials. The Lambda's role does the signing. The user's identity is cryptographically bound to the connection.
Let me quickly go over the full flow from start to finish:
POST /websocket/connect
with JWTnew WebSocket(presignedUrl)
This works well when:
It may not be the right fit when:
The key components:
backend/functions/websocket-connect.js
: The presigned URL generator (the only "middleman")backend/agents/agent/agent.py
: The agent with WebSocket streaming, memory, and toolsfrontend/src/services/websocket.ts
: The browser-side WebSocket clientbackend/template.yaml
: The complete infrastructure definitionI'm really happy with how this turned out. The WebSocket approach removed so much complexity from the architecture and the user experience is noticeably better with real-time streaming. I built this on the side but I've already used the same pattern for several projects at work. The fact that AgentCore handles the WebSocket connection management for us means we don't have to deal with any of the typical WebSocket infrastructure headaches.
Let me know what you think about this approach!
Andres Moreno