# 

# 05-5. WebSocket 基礎

> ⏱️ **閱讀時間：** 18 分鐘
> 🎯 **難度：** ⭐⭐⭐ (進階)

---

## 🤔 一句話解釋

**WebSocket 是雙向通訊協議，讓伺服器可以主動推送資料給客戶端，適合即時應用。**

---

## 🔄 HTTP vs WebSocket

```
HTTP（請求-回應模式）：
┌────────┐         ┌────────┐
│ Client │ ──請求──▶│ Server │
│        │ ◀─回應── │        │
│        │ ──請求──▶│        │
│        │ ◀─回應── │        │
└────────┘         └────────┘
每次通訊都需要新請求

WebSocket（雙向通訊）：
┌────────┐ ═══════ ┌────────┐
│ Client │◀═══════▶│ Server │
│        │  持久    │        │
│        │  連線    │        │
└────────┘ ═══════ └────────┘
一次握手，雙向即時通訊
```

---

## 🎯 使用場景

| 場景 | 說明 |
|------|------|
| **聊天室** | 即時訊息傳遞 |
| **即時通知** | 推播新消息 |
| **股票行情** | 即時價格更新 |
| **遊戲** | 即時狀態同步 |
| **協作編輯** | 多人同時編輯 |

---

## 📦 FastAPI WebSocket 基礎

### 最簡單的範例

```python
from fastapi import FastAPI, WebSocket

app = FastAPI()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 接受連線
    await websocket.accept()

    try:
        while True:
            # 接收訊息
            data = await websocket.receive_text()

            # 回傳訊息
            await websocket.send_text(f"收到: {data}")
    except Exception:
        pass
    finally:
        await websocket.close()
```

### 客戶端測試

```javascript
// 瀏覽器 JavaScript
const ws = new WebSocket("ws://localhost:8000/ws");

ws.onopen = () => {
    console.log("連線成功");
    ws.send("Hello!");
};

ws.onmessage = (event) => {
    console.log("收到:", event.data);
};

ws.onclose = () => {
    console.log("連線關閉");
};
```

---

## 🔧 連線管理

### 連線管理器

```python
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List

app = FastAPI()

class ConnectionManager:
    """WebSocket 連線管理器"""

    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        """接受新連線"""
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        """移除連線"""
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        """發送給單一用戶"""
        await websocket.send_text(message)

    async def broadcast(self, message: str):
        """廣播給所有用戶"""
        for connection in self.active_connections:
            await connection.send_text(message)

manager = ConnectionManager()

@app.websocket("/ws/{client_id}")
async def websocket_endpoint(websocket: WebSocket, client_id: str):
    await manager.connect(websocket)
    await manager.broadcast(f"{client_id} 加入聊天室")

    try:
        while True:
            data = await websocket.receive_text()
            await manager.broadcast(f"{client_id}: {data}")
    except WebSocketDisconnect:
        manager.disconnect(websocket)
        await manager.broadcast(f"{client_id} 離開聊天室")
```

### 房間管理

```python
from collections import defaultdict
from typing import Dict, Set

class RoomManager:
    """聊天室管理器"""

    def __init__(self):
        self.rooms: Dict[str, Set[WebSocket]] = defaultdict(set)
        self.user_rooms: Dict[WebSocket, str] = {}

    async def join_room(
        self,
        websocket: WebSocket,
        room_id: str,
        username: str
    ):
        """加入房間"""
        await websocket.accept()
        self.rooms[room_id].add(websocket)
        self.user_rooms[websocket] = room_id

        await self.broadcast_to_room(
            room_id,
            {"type": "join", "user": username}
        )

    async def leave_room(self, websocket: WebSocket, username: str):
        """離開房間"""
        room_id = self.user_rooms.get(websocket)
        if room_id:
            self.rooms[room_id].discard(websocket)
            del self.user_rooms[websocket]

            await self.broadcast_to_room(
                room_id,
                {"type": "leave", "user": username}
            )

    async def broadcast_to_room(self, room_id: str, message: dict):
        """廣播到房間"""
        import json
        for connection in self.rooms[room_id]:
            await connection.send_text(json.dumps(message))

room_manager = RoomManager()

@app.websocket("/ws/room/{room_id}")
async def room_websocket(
    websocket: WebSocket,
    room_id: str,
    username: str
):
    await room_manager.join_room(websocket, room_id, username)

    try:
        while True:
            data = await websocket.receive_json()
            data["user"] = username
            await room_manager.broadcast_to_room(room_id, data)
    except WebSocketDisconnect:
        await room_manager.leave_room(websocket, username)
```

---

## 🔐 認證

### Token 認證

```python
from fastapi import WebSocket, WebSocketDisconnect, Query, HTTPException
from jose import JWTError, jwt

async def get_current_user_ws(
    websocket: WebSocket,
    token: str = Query(...)
) -> dict:
    """WebSocket 認證"""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        return {"user_id": payload["sub"], "username": payload["name"]}
    except JWTError:
        await websocket.close(code=4001, reason="Invalid token")
        raise HTTPException(status_code=401)

@app.websocket("/ws/chat")
async def authenticated_websocket(
    websocket: WebSocket,
    token: str = Query(...)
):
    # 驗證 token
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        user_id = payload["sub"]
    except JWTError:
        await websocket.close(code=4001, reason="Invalid token")
        return

    await manager.connect(websocket)

    try:
        while True:
            data = await websocket.receive_text()
            # 處理訊息...
    except WebSocketDisconnect:
        manager.disconnect(websocket)
```

### Cookie 認證

```python
from fastapi import Cookie, WebSocket

@app.websocket("/ws/secure")
async def secure_websocket(
    websocket: WebSocket,
    session: str = Cookie(None)
):
    if not session or not verify_session(session):
        await websocket.close(code=4001, reason="Unauthorized")
        return

    await websocket.accept()
    # ...
```

---

## 📊 訊息格式

### JSON 訊息

```python
from pydantic import BaseModel
from typing import Literal
import json

class WSMessage(BaseModel):
    type: Literal["message", "typing", "read"]
    content: str | None = None
    room_id: str | None = None

@app.websocket("/ws/chat/{room_id}")
async def chat_websocket(websocket: WebSocket, room_id: str):
    await websocket.accept()

    try:
        while True:
            # 接收 JSON
            data = await websocket.receive_json()

            # 驗證訊息格式
            try:
                message = WSMessage(**data)
            except Exception:
                await websocket.send_json({
                    "type": "error",
                    "content": "Invalid message format"
                })
                continue

            # 處理不同類型的訊息
            if message.type == "message":
                await manager.broadcast_to_room(room_id, {
                    "type": "message",
                    "content": message.content,
                    "timestamp": datetime.utcnow().isoformat()
                })
            elif message.type == "typing":
                await manager.broadcast_to_room(room_id, {
                    "type": "typing",
                    "user": "username"
                })

    except WebSocketDisconnect:
        pass
```

---

## 💓 心跳機制

```python
import asyncio
from datetime import datetime, timedelta

class HeartbeatManager:
    """心跳管理器"""

    def __init__(self, timeout: int = 30):
        self.timeout = timeout
        self.last_heartbeat: Dict[WebSocket, datetime] = {}

    async def start_heartbeat(self, websocket: WebSocket):
        """開始心跳檢測"""
        self.last_heartbeat[websocket] = datetime.utcnow()

        while True:
            await asyncio.sleep(self.timeout / 2)

            # 檢查是否超時
            last = self.last_heartbeat.get(websocket)
            if not last:
                break

            if datetime.utcnow() - last > timedelta(seconds=self.timeout):
                # 超時，關閉連線
                await websocket.close(code=1000, reason="Heartbeat timeout")
                break

            # 發送 ping
            try:
                await websocket.send_json({"type": "ping"})
            except:
                break

    def update_heartbeat(self, websocket: WebSocket):
        """更新心跳時間"""
        self.last_heartbeat[websocket] = datetime.utcnow()

    def remove(self, websocket: WebSocket):
        """移除連線"""
        self.last_heartbeat.pop(websocket, None)

heartbeat_manager = HeartbeatManager()

@app.websocket("/ws/heartbeat")
async def websocket_with_heartbeat(websocket: WebSocket):
    await websocket.accept()

    # 啟動心跳任務
    heartbeat_task = asyncio.create_task(
        heartbeat_manager.start_heartbeat(websocket)
    )

    try:
        while True:
            data = await websocket.receive_json()

            if data.get("type") == "pong":
                heartbeat_manager.update_heartbeat(websocket)
            else:
                # 處理其他訊息
                pass
    except WebSocketDisconnect:
        pass
    finally:
        heartbeat_task.cancel()
        heartbeat_manager.remove(websocket)
```

---

## 📝 完整聊天室範例

```python
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
from typing import Dict, Set
from datetime import datetime
import json

app = FastAPI()

class ChatMessage(BaseModel):
    type: str
    content: str = ""
    username: str = ""
    timestamp: str = ""

class ChatRoom:
    def __init__(self):
        self.connections: Dict[str, WebSocket] = {}
        self.messages: list = []

    async def connect(self, websocket: WebSocket, username: str):
        await websocket.accept()
        self.connections[username] = websocket

        # 發送歷史訊息
        for msg in self.messages[-50:]:  # 最近 50 條
            await websocket.send_json(msg)

        # 廣播加入訊息
        await self.broadcast({
            "type": "system",
            "content": f"{username} 加入了聊天室",
            "timestamp": datetime.utcnow().isoformat()
        })

    async def disconnect(self, username: str):
        self.connections.pop(username, None)
        await self.broadcast({
            "type": "system",
            "content": f"{username} 離開了聊天室",
            "timestamp": datetime.utcnow().isoformat()
        })

    async def broadcast(self, message: dict):
        self.messages.append(message)

        # 保留最近 100 條訊息
        if len(self.messages) > 100:
            self.messages = self.messages[-100:]

        for connection in self.connections.values():
            try:
                await connection.send_json(message)
            except:
                pass

    async def send_to_user(self, username: str, message: dict):
        ws = self.connections.get(username)
        if ws:
            await ws.send_json(message)

chat_room = ChatRoom()

@app.websocket("/ws/chat")
async def chat_endpoint(websocket: WebSocket, username: str):
    await chat_room.connect(websocket, username)

    try:
        while True:
            data = await websocket.receive_json()

            message = {
                "type": "message",
                "content": data.get("content", ""),
                "username": username,
                "timestamp": datetime.utcnow().isoformat()
            }

            await chat_room.broadcast(message)

    except WebSocketDisconnect:
        await chat_room.disconnect(username)


# HTML 前端
@app.get("/")
async def get_chat_page():
    return HTMLResponse("""
    <!DOCTYPE html>
    <html>
    <head>
        <title>聊天室</title>
        <style>
            #messages { height: 400px; overflow-y: scroll; border: 1px solid #ccc; padding: 10px; }
            .system { color: gray; font-style: italic; }
            .message { margin: 5px 0; }
        </style>
    </head>
    <body>
        <h1>聊天室</h1>
        <div id="messages"></div>
        <input type="text" id="messageInput" placeholder="輸入訊息...">
        <button onclick="sendMessage()">發送</button>

        <script>
            const username = prompt("請輸入你的名稱:");
            const ws = new WebSocket(`ws://localhost:8000/ws/chat?username=${username}`);

            ws.onmessage = (event) => {
                const data = JSON.parse(event.data);
                const messagesDiv = document.getElementById("messages");

                const msgDiv = document.createElement("div");
                msgDiv.className = data.type === "system" ? "system" : "message";

                if (data.type === "message") {
                    msgDiv.innerHTML = `<strong>${data.username}</strong>: ${data.content}`;
                } else {
                    msgDiv.textContent = data.content;
                }

                messagesDiv.appendChild(msgDiv);
                messagesDiv.scrollTop = messagesDiv.scrollHeight;
            };

            function sendMessage() {
                const input = document.getElementById("messageInput");
                if (input.value) {
                    ws.send(JSON.stringify({ content: input.value }));
                    input.value = "";
                }
            }

            document.getElementById("messageInput").addEventListener("keypress", (e) => {
                if (e.key === "Enter") sendMessage();
            });
        </script>
    </body>
    </html>
    """)
```

---

## ✅ 重點總結

### WebSocket 生命週期

| 階段 | 方法 |
|------|------|
| 連線 | `await websocket.accept()` |
| 接收 | `await websocket.receive_text/json()` |
| 發送 | `await websocket.send_text/json()` |
| 關閉 | `await websocket.close()` |

### 最佳實踐

1. 使用連線管理器管理多個連線
2. 實作心跳機制檢測斷線
3. 驗證訊息格式
4. 處理 WebSocketDisconnect 異常
5. 實作認證機制

---

## 🎤 面試這樣答

### Q: WebSocket 和 HTTP 的差別？

**答案：**

> **HTTP：**
> - 單向通訊，客戶端發請求，伺服器回應
> - 每次請求都要建立連線
> - 無狀態
>
> **WebSocket：**
> - 雙向通訊，伺服器可以主動推送
> - 一次握手，持久連線
> - 有狀態，適合即時應用
>
> ```python
> @app.websocket("/ws")
> async def websocket_endpoint(websocket: WebSocket):
>     await websocket.accept()  # 建立連線
>     while True:
>         data = await websocket.receive_text()  # 接收
>         await websocket.send_text(f"Echo: {data}")  # 發送
> ```

---

**上一篇：** [05-4. 並行請求處理](./05-4)
**下一篇：** [05-6. Server-Sent Events](./05-6)

---

最後更新：2025-12-17


---

> 作者: luk  
> URL: https://yoru-karu-blog-lalaluk-52581ac5e0cef170a3c8922c19182ecb6f7bd604.gitlab.io/posts/tutorial/fastapi/05-5/