Skip to content

Cosmos3-edge PR 3: /generate over a persistent WebSocket for control loops - #294

Open
merceod wants to merge 5 commits into
mainfrom
engine/generate-websocket
Open

merceod wants to merge 5 commits into
mainfrom
engine/generate-websocket

Conversation

@merceod

@merceod merceod commented Sep 22, 2026

Copy link
Copy Markdown
Collaborator

Adds /generate/ws, the native /generate request over one persistent WebSocket, for clients that run a control loop against the server (robot policies, streaming world models) and can't afford an HTTP round trip per step.

  • One message = one request with the /generate fields (text, files as [{"name","data"}], input_modalities, output_modalities, model_kwargs, request_id). JSON text frames carry data as base64, msgpack binary frames carry raw bytes, so images and actions skip base64.
  • Replies use the message's encoding, one frame per ResultChunk ({"request_id","modality","data","metadata"}) followed by {"request_id","finish": true}. Errors travel in-band as {"request_id","error"} and the socket stays open.
  • Messages pipeline. Each one gets its own task, replies interleave and carry the request_id. Closing the socket cancels the tasks, and tearing down iter_result_chunks aborts the engine request the same way /generate streaming does.
  • Media are persisted under upload_dir with the same layout as the multipart route (shared _decode_ws_files / _ws_input_layout), so models see identical file_paths / prompt_parts. submit_request and iter_result_chunks are reused unchanged.

Also adds the websockets dependency (uvicorn needs a WS implementation), a section in docs/clients.rst, and examples/cosmos3_action_ws_client.py, an openpi-style policy client that streams observations, decodes float32 [chunk, action_dim] replies and reports chunks/s, actions/s and latency percentiles.

Tests in test/modular/test_generate_ws.py (Starlette in-process client against a fake API server) cover a JSON round trip with media persisted in order, three pipelined msgpack requests, rejected messages answered in-band without dropping the socket, the not-ready close (1013), undecodable / non-object / empty frames, and a mid-stream disconnect cancelling the request. ruff and test/modular pass.

@merceod
merceod force-pushed the engine/generate-websocket branch from e72ae96 to 350fa7c Compare September 23, 2026 23:14

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant