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

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

Ouverte
#2,084 1 commentaire 0 réactions 0 personnes assignées Voir sur GitHub
enhancement Integration SDK
Langage dominant
Python
Étoiles
15.9k
Forks
1.4k
Merge moyen
1 j 17 h
PR mergées (30 j)
163

Description

## 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.

Guide de contribution

Ouvrir le guide de contribution

Évaluation

Cette issue n'a pas encore été évaluée.

Recevez les nouvelles issues par e-mail

Un résumé court des issues GitHub adaptées aux débutants.