FastAPI WebSocket endpoint for real-time updates

<p>I need to push real-time updates to browser clients when traces are validated or new search results arrive. I want to use WebSockets with FastAPI, handle connection cleanup when clients disconnect, and broadcast to multiple connected clients.</p>

<p>Implement a WebSocket manager for broadcast:</p> <div class="highlight"><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">fastapi</span><span class="w"> </span><span class="kn">import</span> <span class="n">FastAPI</span><span class="p">,</span> <span class="n">WebSocket</span><span class="p">,</span> <span class="n">WebSocketDisconnect</span> <span class="kn">from</span><span class="w"> </span><span class="nn">typing</span><span class="w"> </span><span class="kn">import</span> <span class="n">Any</span> <span class="kn">import</span><span class="w"> </span><span class="nn">json</span> <span class="k">class</span><span class="w"> </span><span class="nc">ConnectionManager</span><span class="p">:</span> <span class="k">def</span><span class="w"> </span><span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="n">WebSocket</span><span class="p">]</span> <span class="o">=</span> <span class="p">[]</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">connect</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ws</span><span class="p">:</span> <span class="n">WebSocket</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="k">await</span> <span class="n">ws</span><span class="o">.</span><span class="n">accept</span><span class="p">()</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="k">def</span><span class="w"> </span><span class="nf">disconnect</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">ws</span><span class="p">:</span> <span class="n">WebSocket</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">broadcast</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">message</span><span class="p">:</span> <span class="n">Any</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="n">data</span> <span class="o">=</span> <span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">(</span><span class="n">message</span><span class="p">)</span> <span class="n">disconnected</span> <span class="o">=</span> <span class="p">[]</span> <span class="k">for</span> <span class="n">ws</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections</span><span class="p">:</span> <span class="k">try</span><span class="p">:</span> <span class="k">await</span> <span class="n">ws</span><span class="o">.</span><span class="n">send_text</span><span class="p">(</span><span class="n">data</span><span class="p">)</span> <span class="k">except</span> <span class="ne">Exception</span><span class="p">:</span> <span class="n">disconnected</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="k">for</span> <span class="n">ws</span> <span class="ow">in</span> <span class="n">disconnected</span><span class="p">:</span> <span class="bp">self</span><span class="o">.</span><span class="n">connections</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="n">manager</span> <span class="o">=</span> <span class="n">ConnectionManager</span><span class="p">()</span> <span class="nd">@app</span><span class="o">.</span><span class="n">websocket</span><span class="p">(</span><span class="s1">'/ws/traces'</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">traces_websocket</span><span class="p">(</span><span class="n">ws</span><span class="p">:</span> <span class="n">WebSocket</span><span class="p">):</span> <span class="k">await</span> <span class="n">manager</span><span class="o">.</span><span class="n">connect</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="k">try</span><span class="p">:</span> <span class="k">while</span> <span class="kc">True</span><span class="p">:</span> <span class="c1"># Keep connection alive — wait for client messages or ping</span> <span class="k">await</span> <span class="n">ws</span><span class="o">.</span><span class="n">receive_text</span><span class="p">()</span> <span class="c1"># or receive_json()</span> <span class="k">except</span> <span class="n">WebSocketDisconnect</span><span class="p">:</span> <span class="n">manager</span><span class="o">.</span><span class="n">disconnect</span><span class="p">(</span><span class="n">ws</span><span class="p">)</span> <span class="c1"># Broadcast from any route:</span> <span class="nd">@router</span><span class="o">.</span><span class="n">post</span><span class="p">(</span><span class="s1">'/traces'</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">create_trace</span><span class="p">(</span><span class="n">body</span><span class="p">:</span> <span class="n">TraceCreate</span><span class="p">):</span> <span class="n">trace</span> <span class="o">=</span> <span class="k">await</span> <span class="n">save_trace</span><span class="p">(</span><span class="n">body</span><span class="p">)</span> <span class="k">await</span> <span class="n">manager</span><span class="o">.</span><span class="n">broadcast</span><span class="p">({</span><span class="s1">'type'</span><span class="p">:</span> <span class="s1">'trace_created'</span><span class="p">,</span> <span class="s1">'id'</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">trace</span><span class="o">.</span><span class="n">id</span><span class="p">)})</span> <span class="k">return</span> <span class="n">trace</span> </code></pre></div> <p>Key points: - Always handle <code>WebSocketDisconnect</code> — clients can disconnect at any time - Track disconnected sockets during broadcast — don't modify list while iterating - For production scale, use Redis pub/sub as the broadcast backend - Consider authentication: check API key in query param or first message</p>