feat: add AG-UI protocol support via serve_ag_ui and AGUIApp

[tejaskash]

Summary

• Adds AGUIApp, serve_ag_ui(), and build_ag_ui_app() for AG-UI protocol support
• Single @app.entrypoint handler is served on both transports per the AG-UI contract:
  • POST /invocations — SSE (unidirectional streaming)
  • /ws — WebSocket (bidirectional, same AG-UI events)
• Bedrock header extraction, Docker host detection, /ping health check
• Pre-stream errors return HTTP 400 / WS close 1003; mid-stream errors emit RunErrorEvent
• ag-ui-protocol>=0.1.10 added as optional dependency (pip install "bedrock-agentcore[ag-ui]")
• Follows existing a2a.py patterns (lazy imports, _check_sdk(), _takes_context, _build_request_context)

Sample: Framework agent (e.g. Strands + ag-ui-strands)

For agents that already expose a .run() method (Strands, LangGraph adapters, etc.):

from strands import Agent
from strands.models.bedrock import BedrockModel
from ag_ui_strands import StrandsAgent
from bedrock_agentcore.runtime import serve_ag_ui

model = BedrockModel(model_id="us.anthropic.claude-sonnet-4-20250514-v1:0")
strands_agent = Agent(model=model, system_prompt="You are a helpful assistant.")

agui_agent = StrandsAgent(
    agent=strands_agent,
    name="my_agent",
    description="A helpful assistant",
)

# One line — serves on /invocations (SSE), /ws (WebSocket), and /ping
serve_ag_ui(agui_agent)

Sample: Custom agent (decorator form)

For writing agent logic directly without a framework adapter:

from bedrock_agentcore.runtime import AGUIApp
from bedrock_agentcore.runtime.context import BedrockAgentCoreContext, RequestContext
from ag_ui.core import (
    RunAgentInput,
    RunStartedEvent,
    RunFinishedEvent,
    TextMessageStartEvent,
    TextMessageContentEvent,
    TextMessageEndEvent,
)

app = AGUIApp()

@app.entrypoint
async def my_agent(input_data: RunAgentInput, context: RequestContext):
    # Access Bedrock headers
    token = BedrockAgentCoreContext.get_workload_access_token()

    yield RunStartedEvent(thread_id=input_data.thread_id, run_id=input_data.run_id)

    # Stream a text message
    msg_id = "msg-1"
    yield TextMessageStartEvent(message_id=msg_id, role="assistant")

    # Call your LLM, yield chunks as they arrive
    for chunk in ["Hello", ", ", "world!"]:
        yield TextMessageContentEvent(message_id=msg_id, delta=chunk)

    yield TextMessageEndEvent(message_id=msg_id)
    yield RunFinishedEvent(thread_id=input_data.thread_id, run_id=input_data.run_id)

@app.ping
def health():
    from bedrock_agentcore.runtime.models import PingStatus
    return PingStatus.HEALTHY

app.run()  # Serves /invocations (SSE), /ws (WebSocket), /ping on port 8080

Sample: Testing with TestClient

from bedrock_agentcore.runtime import build_ag_ui_app
from starlette.testclient import TestClient

app = build_ag_ui_app(my_agent)
client = TestClient(app)

# Test SSE transport
resp = client.post("/invocations", json={...})
assert resp.status_code == 200

# Test WebSocket transport
with client.websocket_connect("/ws") as ws:
    ws.send_json({...})
    msg = ws.receive_text()  # AG-UI event

Files

Action	Path
Create	src/bedrock_agentcore/runtime/ag_ui.py — AGUIApp, serve_ag_ui, build_ag_ui_app
Modify	src/bedrock_agentcore/runtime/__init__.py — export new symbols
Modify	pyproject.toml — add ag-ui optional dep group + dev dep
Modify	uv.lock — updated lockfile
Modify	README.md — AG-UI section with examples
Create	docs/proposals/ag-ui-protocol-support.md — updated design doc
Create	tests/bedrock_agentcore/runtime/test_ag_ui.py — 20 unit tests
Create	tests/integration/runtime/test_ag_ui_integration.py — 11 integration tests

Test plan

• pytest tests/bedrock_agentcore/runtime/test_ag_ui.py — 20 unit tests pass
• pytest tests/integration/runtime/test_ag_ui_integration.py — 11 integration tests pass
• pytest tests/bedrock_agentcore/runtime/ — all 252 runtime tests pass (no regressions)
• ruff check + ruff format --check — clean

Manual testing

Verified end-to-end with four framework integrations and both AG-UI transports (SSE and WebSocket).

1. Strands on Bedrock (SSE)

Scaffolded a sample project using ag-ui-strands + StrandsAgent wrapping a Strands Agent with BedrockModel (Claude Sonnet). Served via serve_ag_ui(agui_agent). Connected a CopilotKit React frontend (HttpAgent → POST /invocations). Confirmed streaming text responses over SSE.

2. Google ADK (WebSocket)

Replaced Strands with google-adk using Agent + Runner + InMemorySessionService. Used the @app.entrypoint decorator form to map ADK runner events to AG-UI events. Connected a custom React chat UI via WebSocket to /ws. Confirmed bidirectional WebSocket transport — sent RunAgentInput as first JSON frame, received AG-UI events as text frames, connection closed on completion.

3. LangChain on Bedrock (SSE)

Replaced ADK with langchain-aws (ChatBedrockConverse). Used @app.entrypoint decorator with model.astream() to stream LangChain chunks as TextMessageContentEvents. Connected CopilotKit React frontend (HttpAgent → POST /invocations). Confirmed streaming text responses over SSE.

4. LangGraph on Bedrock (SSE) — native AG-UI integration

Used LangGraphAGUIAgent from copilotkit wrapping a compiled LangGraph StateGraph (with MemorySaver checkpointer) and ChatBedrockConverse (Claude Sonnet). The native ag-ui-langgraph adapter handles all AG-UI event translation — no manual event yielding needed. Passed the agent directly to serve_ag_ui(agui_agent), which detected the .run() method and wired it to both transports. Connected a CopilotKit React frontend (HttpAgent → POST /invocations) with ExperimentalEmptyAdapter. Confirmed streaming text responses over SSE end-to-end.

# Working LangGraph sample — serve_ag_ui with native AG-UI adapter
from copilotkit import LangGraphAGUIAgent
from langchain_aws import ChatBedrockConverse
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import MessagesState, StateGraph
from bedrock_agentcore.runtime import serve_ag_ui

model = ChatBedrockConverse(model="us.anthropic.claude-sonnet-4-20250514-v1:0", region_name="us-west-2")

async def chat(state: MessagesState):
    response = await model.ainvoke(state["messages"])
    return {"messages": [response]}

graph = StateGraph(MessagesState)
graph.add_node("chat", chat)
graph.set_entry_point("chat")
graph.set_finish_point("chat")
compiled_graph = graph.compile(checkpointer=MemorySaver())

serve_ag_ui(LangGraphAGUIAgent(name="agent", description="Assistant", graph=compiled_graph))

Transports tested

Transport	Endpoint	Tested with
SSE	POST /invocations	Strands, LangChain, LangGraph, curl
WebSocket	/ws	ADK, custom React WebSocket client
Health	GET /ping	curl, TestClient

curl verification

# SSE
curl -N -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"thread_id":"t1","run_id":"r1","state":[],"messages":[{"role":"user","content":"What is the capital of France?","id":"m1"}],"tools":[],"context":[],"forwardedProps":{}}'

# Ping
curl http://localhost:8080/ping

Both returned expected AG-UI event stream and {"status":"Healthy",...} respectively.

[sundargthb]

Two things to verify against the ag-ui-protocol SDK before merging:

The curl example uses thread_id/run_id (snake_case), but the AG-UI spec defines these as threadId/runId (camelCase). Confirm that RunAgentInput(**payload) handle both formats.

The error path uses RunErrorEvent(message=str(e)) — does RunErrorEvent from ag_ui.core accept message as a constructor parameter? The AG-UI spec defines RUN_ERROR with both code and message fields. Passing only message might be fine as a default, but confirm that the constructor signature and whether we should include a code (e.g. "INTERNAL_ERROR") for spec compliance.​​​​​​​​​​​​​​​​

Otherwise, LGTM

[tejaskash]

Thanks for the review! Verified both points against ag-ui-protocol:

1. snake_case vs camelCase in RunAgentInput

RunAgentInput is a Pydantic model with populate_by_name=True (alias support), so it accepts both formats:

# camelCase (AG-UI spec / what CopilotKit sends)
RunAgentInput(**{"threadId": "t1", "runId": "r1", ...})  # ✅

# snake_case (curl convenience / Python convention)
RunAgentInput(**{"thread_id": "t1", "run_id": "r1", ...})  # ✅

Both work. The curl example uses snake_case for readability, but real AG-UI clients (CopilotKit, etc.) send camelCase — both are handled.

2. RunErrorEvent constructor signature

RunErrorEvent(*, message: str, code: Optional[str] = None, ...)

• message is required, code is optional (defaults to None)
• Our usage RunErrorEvent(message=str(e)) is valid per the SDK

That said, adding code="INTERNAL_ERROR" for mid-stream errors would be more spec-compliant. I'll add that.