Skip to content

feat: pluggable client transport with a plain WebSocket option - #83

Merged
mm-zacharydavison merged 4 commits into
mainfrom
feat/pluggable-transport
Sep 11, 2026
Merged

mm-zacharydavison merged 4 commits into
mainfrom
feat/pluggable-transport

Conversation

@mm-zacharydavison

@mm-zacharydavison mm-zacharydavison commented Sep 3, 2026 •

Copy link
Copy Markdown
Contributor

🧑‍🦱 @masudahiroto & @yuma2024 : I want to use use-ai for Cordless UI, so we can also be users of the library and contribute / test it.
But Cordless is a Go application, so cannot run Socket.IO for transport, so I added support for plain WebSocket transport, and an abstraction for transport.

Please check it.

Summary

UseAIClient builds its own Socket.IO socket today, so a Socket.IO server is the only server it can reach. This PR gives the client a transport seam and ships two transports.

  • SocketIOTransport keeps today's behaviour. It is the default.
  • WebSocketTransport carries AG-UI events as JSON text frames over a plain WebSocket. Use it to reach a server that does not serve Socket.IO. That server does not have to be Node.

Nothing existing changes. new UseAIClient(url) still works. A UseAIProvider given serverUrl still connects over Socket.IO.

Client

  • UseAITransport has five members: connect, disconnect, send, on, connected.
  • UseAIClient takes string | UseAITransport.
  • UseAIConfig is a discriminated union: give UseAIProvider exactly one of serverUrl or transport. The provider reads transport once, on the first render, so an inline object does not reconnect the client on every render. UseAIContextValue.serverUrl becomes optional.
  • WebSocketTransport reconnects through partysocket, with the same limits as Socket.IO: indefinite, one second doubling to ten.
  • The transport-upgrade logging and connect_error stay inside SocketIOTransport. Both are specific to Socket.IO.

Server

  • transport?: 'socketio' | 'websocket' (default 'socketio'). The server runs one listener. With 'websocket', it does not create a Socket.IO server; it accepts WebSocket upgrades at /.
  • Both runtime adapters take a RuntimeListener union. Bun uses Bun.serve's own upgrade. Node uses ws in noServer mode.
  • ClientSession.socket narrows from Socket to ClientConnection (id, connected, emit). A Socket.IO socket satisfies it structurally, so FeedbackPlugin and WorkflowsPlugin are unchanged.
  • The Docker image reads TRANSPORT.

Framing

Upstream, the client sends each UseAIClientMessage as one frame, with nothing around it. Downstream, the server sends one AG-UI event per frame. agents and config travel as AG-UI CUSTOM events. The client ignores an event type or a custom name it does not know. See docs/websocket-protocol.md.

Tests

  • client.test.ts runs table-driven over both transports: 64 tests, up from 29. Each harness drives its transport at the wire level.
  • SocketIOTransport.test.ts keeps the socket.io-client module mock and covers connect_error, the io() options and upgrade logging.
  • WebSocketTransport.test.ts covers framing, CUSTOM routing, unknown names, malformed frames, connected, reconnection, and disconnect() stopping it.
  • websocket-transport.integration.test.ts drives a real WebSocketTransport against a real transport: 'websocket' server through prompt → tool call → tool result → RUN_FINISHED, on both the Bun and Node adapters. It also checks that a default server refuses a plain WebSocket.

bun run test and bun run build pass. apps/example passes suggestions.e2e with real API calls. tool-approval.e2e and dynamic-tool-registration.e2e fail identically on main, so they are pre-existing.

Dependencies

packages/client adds partysocket. packages/server adds ws as a dependency and @meetsmore-oss/use-ai-client as a devDependency, so the integration test imports the real transport rather than a stub.

Not in this PR

No version bump or release. The change is additive, so a minor bump fits when you release.

🤖 Generated with Claude Code

UseAIClient now takes a UseAITransport instead of building its own
Socket.IO socket. Two transports ship: SocketIOTransport keeps today's
behaviour and stays the default, WebSocketTransport carries JSON frames
over a plain WebSocket for servers that do not speak Socket.IO.

The server serves both on one port. webSocketPath (default '/ws') adds a
plain listener beside Socket.IO in both runtime adapters. ClientSession.socket
narrows to ClientConnection, which a Socket.IO socket satisfies, so plugins
are unchanged.

UseAIProvider accepts a transport prop, read once so an inline object does
not churn the connection. new UseAIClient(url) and a provider given only
serverUrl behave exactly as before.

The UseAIClient suite now runs table-driven over both transports. An
integration test drives a real WebSocketTransport through a full run on
both the Bun and Node adapters.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@mm-zacharydavison
mm-zacharydavison marked this pull request as draft September 3, 2026 16:09
Comment thread docs/websocket-protocol.md
Comment thread docs/websocket-protocol.md Outdated
Comment thread docs/websocket-protocol.md Outdated
Comment thread docs/websocket-protocol.md
Comment thread packages/client/src/transport/WebSocketTransport.ts
mm-zacharydavison and others added 3 commits September 4, 2026 11:54
…cket

- UseAIProvider takes exactly one of serverUrl or transport, as a
  discriminated union on UseAIConfig. Context serverUrl becomes optional.
- The server runs one listener: transport: 'socketio' | 'websocket',
  default 'socketio'. No webSocketPath. The plain listener serves at '/'.
  The Docker image reads TRANSPORT.
- Downstream frames are bare AG-UI events. agents and config travel as
  AG-UI CUSTOM events, so the wire is AG-UI rather than a bespoke envelope.
- WebSocketTransport reconnects through partysocket instead of its own loop.
- Protocol doc rewritten.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
UseAITransport gains a readonly url. Both bundled transports already take
one. UseAIContextValue.serverUrl is a required string again, populated from
the transport when the provider was given one.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Client: UseAITransport is now onEvent(AGUIEvent) + onConnectionChange.
The five named channels, the handler registry and the agents/config
special case in WebSocketTransport are gone. SocketIOTransport presents
its legacy agents/config events as AG-UI CUSTOM events, and
UseAIClient.handleEvent gains one CUSTOM branch. SocketIOTransport loses
options nothing used. The provider resolves its target once.

Server: ClientConnection is the whole per-connection boundary (id,
ipAddress, connected, emit, onMessage, onClose). SocketIOClientConnection
and WebSocketClientConnection own their protocol details; server.ts has
one acceptConnection. handleClientMessage takes the session it is given.
The agents payload is built once. Bun adapter returns {path, upgrade,
websocket} per listener and keeps the connection on ws.data. Binary
frames are ignored on both runtimes. Listener member types are named.
X-Forwarded-For parsing is one helper.

Tests share one FakeWebSocket, one socket.io mock and one waitUntil.
The negative transport test checks a raw socket close instead of
waiting out a 5s timeout.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@mm-zacharydavison

Copy link
Copy Markdown
Contributor Author

Simplification pass in 6970d43, net -271 lines on the branch.

  • UseAITransport is now onEvent(AGUIEvent) + onConnectionChange. The five Socket.IO-shaped channels, handlerRegistry.ts and UseAITransportEventName are gone. SocketIOTransport presents its legacy agents/config events as AG-UI CUSTOM, so a third-party transport only has to deliver AG-UI events.
  • Server: ClientConnection is the whole per-connection boundary. SocketIOClientConnection / WebSocketClientConnection own their protocol details and server.ts has one acceptConnection.
  • Bun adapter: { path, upgrade, websocket } per listener, connection kept on ws.data, no tri-state return, no WeakMap. Named SocketIOListener / WebSocketListener types replace the infer gymnastics.
  • Dropped SocketIOTransportOptions (nothing used it), the unused client.serverUrl getter, the duplicated test fakes, the 5 s negative test, and the per-message session lookup.

@mm-zacharydavison
mm-zacharydavison requested review from masudahiroto and yuma2024 and removed request for masudahiroto September 4, 2026 10:40
@mm-zacharydavison
mm-zacharydavison marked this pull request as ready for review September 4, 2026 10:40
@mm-zacharydavison
mm-zacharydavison merged commit b5fb2a6 into main Sep 11, 2026
1 check passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Development

Successfully merging this pull request may close these issues.

2 participants