ag-ui-protocol / ag-ui-protocol/ag-ui

feat(transport): WebSocket support for ag-ui-adk and @ag-ui/client

Abierto
#2,084 1 comentario 0 reacciones 0 asignados Ver en GitHub
enhancement Integration SDK
Lenguaje dominante
Python
Estrellas
15.9k
Forks
1.4k
Merge medio
1 d 17 h
PR fusionados (30 d)
163

Descripción

## Problem

AG-UI currently only supports SSE (Server-Sent Events) as a transport. While SSE works well, it has friction with certain enterprise infrastructure:

- **Envoy / API gateways**: SSE requires `timeout: 0s` and `X-Accel-Buffering: no` configuration; some proxies buffer or drop long-lived HTTP/1.1 responses.
- **HITL (Human-in-the-Loop)**: The current SSE model requires *two separate HTTP connections* per HITL interaction — one that pauses on `RUN_FINISHED` (no `TOOL_CALL_RESULT`) and a second POST to resume. With WebSocket the connection stays open and the client sends the `ToolMessage` over the same channel, which is more natural.
- **Bidirectional use-cases**: Future multi-turn streaming and live tool result submission are natural fits for WebSocket.

The protocol already models this: `TransportCapabilities.websocket` exists in `@ag-ui/core` and `capabilities.py`, but no integration implements it.

## Proposal

Add WebSocket transport support to:

1. **`integrations/adk-middleware/python`** — a `ws_endpoint` alongside the existing SSE endpoint in `endpoint.py`, using FastAPI's native `WebSocket` support. The `ADKAgent` internals (event queue, translator) are unchanged; only the output channel switches from SSE to WebSocket frames.

2. **`sdks/typescript/packages/client`** — a `WebSocketAgent` class extending `AbstractAgent`. Wire-format: one JSON frame per AG-UI event, same schema as SSE today.

## Scope

- No changes to the AG-UI event schema, `RunAgentInput`, or any other integration.
- The SSE path is fully preserved; WebSocket is **additive**.
- `capabilities_endpoint` updated to advertise `transport.websocket: true` when the WS endpoint is mounted.
- HITL over WebSocket: the client sends the `ToolMessage` frame over the same connection instead of a new POST.

## Backend sketch (`endpoint.py`)

```python
@app.websocket("/ws")
async def ws_adk_endpoint(websocket: WebSocket):
await websocket.accept()
try:
data = await websocket.receive_json()
input_data = RunAgentInput(**data)
async for event in _run_agent(agent, input_data):
await websocket.send_text(
event.model_dump_json(by_alias=True, exclude_none=True)
)
except WebSocketDisconnect:
pass
```

## Frontend sketch (`ws-agent.ts` in `@ag-ui/client`)

```typescript
export class WebSocketAgent extends AbstractAgent {
constructor(private readonly url: string, config?: AgentConfig) {
super(config);
}

run(input: RunAgentInput): Observable {
return new Observable(subscriber => {
const ws = new WebSocket(this.url);
ws.onopen = () => ws.send(JSON.stringify(input));
ws.onmessage = ({ data }) => {
const event = parseEvent(JSON.parse(data));
subscriber.next(event);
if (event.type === EventType.RUN_FINISHED ||
event.type === EventType.RUN_ERROR) {
subscriber.complete();
ws.close();
}
};
ws.onerror = e => subscriber.error(e);
ws.onclose = () => subscriber.complete();
return () => ws.close();
});
}
}
```

## Motivation

Surfaced by an enterprise user integrating ADK agents behind an Envoy API gateway. The existing SSE endpoint works but requires non-obvious proxy configuration (`timeout: 0s`, `X-Accel-Buffering: no`). WebSocket is a cleaner fit for their infrastructure and enables simpler HITL flows without a second HTTP connection.

`TransportCapabilities` in `@ag-ui/core` already reserves `websocket: boolean` for exactly this purpose — this PR would be the first implementation of it.

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.