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

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

Offen
#2,084 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
enhancement Integration SDK
Vorherrschende Sprache
Python
Sterne
15.9k
Forks
1.4k
Ø Merge
1 T. 17 Std.
Gemergte PRs (30 T.)
163

Beschreibung

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

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.