# How the WebSocket Communication Layer Streams Code Previews in screenshot-to-code

> Discover how the WebSocket communication layer streams code previews for screenshot-to-code. Learn about the /generate-code endpoint and its JSON message types like setCode and status.

- Repository: [Abi Raja/screenshot-to-code](https://github.com/abi/screenshot-to-code)
- Tags: internals
- Published: 2026-03-02

---

**The screenshot-to-code backend exposes a `/generate-code` WebSocket endpoint that drives a middleware pipeline to stream LLM-generated code previews via JSON messages with types like `setCode`, `status`, and `variantComplete`.**

The abi/screenshot-to-code repository transforms UI screenshots into functional HTML using large language models. Its WebSocket communication layer enables real-time collaboration by streaming incremental code previews to the frontend through a structured pipeline of middleware components and a type-safe JSON protocol.

## WebSocket Endpoint and Connection Setup

In [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) (lines 10-14), the `stream_code` function establishes the WebSocket connection at the `/generate-code` endpoint. When a client connects, the function initializes a `Pipeline` instance and executes it to process the request asynchronously.

The pipeline begins with **WebSocketSetupMiddleware**, which creates a `WebSocketCommunicator` instance, accepts the connection, and guarantees proper cleanup by invoking `close()` when the pipeline terminates, regardless of success or failure.

## Middleware Pipeline Architecture for Streaming Code Previews

The streaming architecture follows a chain-of-responsibility pattern where each middleware handles a specific concern. The pipeline processes requests through six distinct stages defined in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py):

- **WebSocketSetupMiddleware**: Manages connection lifecycle and instantiates the communicator
- **ParameterExtractionMiddleware**: Validates incoming JSON payloads via `receive_params()` (lines 72-80)
- **StatusBroadcastMiddleware**: Pushes initial status messages to the UI before generation begins (lines 90-112)
- **PromptCreationMiddleware**: Builds LLM prompt lists for internal state management
- **CodeGenerationMiddleware**: Orchestrates concurrent variant generation through `AgenticGenerationStage` (lines 136-168)
- **PostProcessingMiddleware**: Handles logging without sending additional UI messages

### WebSocketSetupMiddleware and Communicator Management

This middleware instantiates the `WebSocketCommunicator` class defined in [`backend/ws/__init__.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/__init__.py). It wraps the raw WebSocket and provides the `send_message` method used throughout the pipeline. According to the source code (lines 49-63), this middleware ensures the socket closes with custom code `4332`—defined in [`backend/ws/constants.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/constants.py)—when processing completes.

### Parameter Extraction and Validation

After connection establishment, **ParameterExtractionMiddleware** receives the JSON payload, validates the request structure, and stores parameters in the pipeline context for downstream access. This stage ensures all required generation parameters are present before invoking expensive LLM operations.

### Status Broadcasting for UI Initialization

Before generation begins, **StatusBroadcastMiddleware** informs the frontend how many code variants will run. It pushes initial `status` messages like "Generating code…" for each variant (lines 90-112), allowing the UI to prepare placeholder containers for parallel previews and establish progress tracking.

### Code Generation and Real-Time Streaming

The **CodeGenerationMiddleware** selects appropriate LLM models and executes each variant through `AgenticGenerationStage`. Located at lines 136-168 in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py), this middleware handles the actual streaming of generated code back to the client as the LLM produces output, enabling near-real-time preview updates.

## JSON Message Protocol for Code Preview Streaming

All communication flows through `WebSocketCommunicator.send_message` in [`backend/ws/__init__.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/__init__.py) (lines 65-95), which constructs JSON payloads and transmits them via `websocket.send_json`. Every message includes a `type` field determining frontend behavior and a `variantIndex` field (0-based) identifying which preview panel should update.

### Message Types and Their Purposes

The protocol defines distinct message types for different stages of the generation process:

- **`setCode`**: Transmits the complete HTML code for a specific variant. Sent after `AgenticGenerationStage._run_variant` finishes processing at line 90.
- **`status`**: Delivers human-readable progress updates such as "Generating code…" or "thinking". Originating from `StatusBroadcastMiddleware` and the agent itself during generation.
- **`variantComplete`**: Signals successful completion of a variant's generation cycle, sent immediately after `setCode` at line 92.
- **`variantError`**: Communicates errors specific to one variant when exceptions occur in `_run_variant` (lines 100-124).
- **`error`**: Indicates fatal errors that terminate the entire connection via `WebSocketCommunicator.throw_error` (lines 100-108).

## Connection Lifecycle and Error Handling

The pipeline manages connection state through structured error boundaries. When `AgenticGenerationStage` encounters exceptions during variant generation, it catches them and transmits `variantError` messages without terminating other concurrent variants.

For unrecoverable failures, `WebSocketCommunicator.throw_error` sends an `error` type message and initiates connection closure. The custom close code `4332` defined in [`backend/ws/constants.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/constants.py) allows the frontend to distinguish between normal completion and error-induced termination.

## Implementation Examples

### Server-Side: Streaming Generated Code

Inside `AgenticGenerationStage._run_variant`, the backend streams results as they become available:

```python

# Inside AgenticGenerationStage._run_variant

completion = await runner.run(model, prompt_messages)
if completion:
    await self.send_message("setCode", completion, index, None, None)
await self.send_message(
    "variantComplete",
    "Variant generation complete",
    index,
    None,
    None,
)

```

### Server-Side: Low-Level Message Transmission

The `WebSocketCommunicator` class in [`backend/ws/__init__.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/__init__.py) implements the actual JSON serialization:

```python
async def send_message(
    self,
    type: MessageType,
    value: str | None,
    variantIndex: int,
    data: Dict[str, Any] | None = None,
    eventId: str | None = None,
) -> None:
    payload = {"type": type, "variantIndex": variantIndex}
    if value is not None:
        payload["value"] = value
    if data is not None:
        payload["data"] = data
    if eventId is not None:
        payload["eventId"] = eventId
    await self.websocket.send_json(payload)

```

### Client-Side: Consuming the Stream

The frontend establishes the connection and routes messages based on type, as implemented in components referencing [`frontend/src/types.ts`](https://github.com/abi/screenshot-to-code/blob/main/frontend/src/types.ts):

```ts
const ws = new WebSocket(`wss://${location.host}/generate-code`);
ws.onmessage = (event) => {
  const msg = JSON.parse(event.data);
  switch (msg.type) {
    case "setCode":
      previewPanels[msg.variantIndex].innerHTML = msg.value;
      break;
    case "status":
      statusBars[msg.variantIndex].textContent = msg.value;
      break;
    case "variantComplete":
      console.log(`Variant ${msg.variantIndex} finished`);
      break;
    case "variantError":
      console.error(`Variant ${msg.variantIndex} failed:`, msg.value);
      break;
  }
};

```

## Summary

- The `/generate-code` endpoint in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py) initiates a middleware pipeline that orchestrates WebSocket communication.
- **WebSocketSetupMiddleware** manages connection lifecycle using `WebSocketCommunicator` from [`backend/ws/__init__.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/__init__.py) and closes with code `4332` from [`backend/ws/constants.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/constants.py).
- Six middleware stages process parameters, broadcast initial status, and execute LLM generation through `AgenticGenerationStage`.
- The JSON protocol uses `type` fields (`setCode`, `status`, `variantComplete`, `variantError`, `error`) and `variantIndex` to route updates to specific UI panels.
- Errors are isolated per variant via `variantError` messages, while fatal errors use the `error` type before socket termination.
- Frontend implementations listen for incoming JSON messages and update preview panels in real-time as code generation progresses.

## Frequently Asked Questions

### How does the backend handle multiple code variants simultaneously?

The `CodeGenerationMiddleware` runs multiple variants concurrently through `AgenticGenerationStage._run_variant` at lines 136-168 in [`backend/routes/generate_code.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/routes/generate_code.py). Each variant receives a unique `variantIndex` (0-based), and all status and code messages include this index so the frontend can update the correct preview panel independently without blocking other generations.

### What happens when a single variant fails but others are still running?

The backend catches exceptions within `_run_variant` (lines 100-124) and transmits a `variantError` message specific to the failed index. This allows the frontend to display error states for individual variants while continuing to receive progress updates and final code from successfully completing variants.

### Can the frontend distinguish between a completed generation and a connection error?

Yes. Normal pipeline completion triggers `WebSocketCommunicator.close()` with custom close code `4332` defined in [`backend/ws/constants.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/constants.py). Fatal errors use the `error` message type via `throw_error` before closing the socket, allowing the frontend to differentiate between successful termination and error states through both message types and close codes.

### Where is the message type schema defined for type safety?

The backend uses Python type hints with `MessageType` literals in [`backend/ws/__init__.py`](https://github.com/abi/screenshot-to-code/blob/main/backend/ws/__init__.py). The frontend mirrors this contract in [`frontend/src/types.ts`](https://github.com/abi/screenshot-to-code/blob/main/frontend/src/types.ts) through the `MessageType` union, ensuring both ends agree on valid message types like `setCode`, `status`, and `variantComplete` for reliable TypeScript and Python integration.