Skip to content

own websockets - own connection tracking fixed #9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
DoktorShift merged 1 commit into from
Feb 2, 2026
Merged

own websockets - own connection tracking fixed #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 4 additions & 5 deletions static/js/index.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -410,8 +410,8 @@ window.app = Vue.createApp({
return
}

// Use extension WebSocket endpoint
const websocketUrl = this.wsLocation + '/devicetimer/api/v1/ws/' + deviceId
// Use extension WebSocket endpoint with browser type (doesn't count as hardware connection)
const websocketUrl = this.wsLocation + '/devicetimer/api/v1/ws/' + deviceId + '?type=browser'
this.websocketMessage = 'Connecting...'
this.activeWebsocketDeviceId = deviceId

Expand All @@ -420,8 +420,7 @@ window.app = Vue.createApp({
this.activeWebsocket = ws

ws.onopen = () => {
this.websocketMessage = 'Connected'
this.fetchConnectionStatus()
this.websocketMessage = 'Watching for payments...'
}

ws.onmessage = (event) => {
Expand All @@ -431,7 +430,7 @@ window.app = Vue.createApp({
}

ws.onclose = () => {
this.websocketMessage = 'Disconnected'
this.websocketMessage = ''
if (this.activeWebsocket === ws) {
this.activeWebsocket = null
this.activeWebsocketDeviceId = null
Expand Down
6 changes: 3 additions & 3 deletions templates/devicetimer/index.html
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1096,12 +1096,12 @@
</div>

<!-- WebSocket Status -->
<div class="q-mt-sm">
<div class="q-mt-sm" v-if="websocketMessage">
<q-chip
:color="websocketMessage.includes('connected') ? 'positive' : websocketMessage.includes('Payment') ? 'positive' : 'grey-6'"
:color="websocketMessage.includes('Payment') ? 'positive' : 'grey-6'"
text-color="white"
size="sm"
:icon="websocketMessage.includes('connected') ? 'wifi' : websocketMessage.includes('Payment') ? 'check_circle' : 'wifi_off'"
:icon="websocketMessage.includes('Payment') ? 'check_circle' : 'visibility'"
>
{% raw %}{{ wsMessage }}{% endraw %}
</q-chip>
Expand Down
117 changes: 70 additions & 47 deletions websocket.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,93 +2,116 @@
WebSocket management for DeviceTimer extension.

Handles device connections and message broadcasting.
Tracks connected devices to show real-time status in UI.
Tracks connected hardware devices to show real-time status in UI.
Browser connections (for watching payments) are tracked separately.
"""

from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Query
from loguru import logger
from typing import Dict, Set
from typing import Dict, Set, Optional

devicetimer_websocket_router = APIRouter()

# Track connected devices: device_id -> set of WebSocket connections
# Multiple connections per device are supported (e.g., multiple browser tabs)
_connected_clients: Dict[str, Set[WebSocket]] = {}
# Track hardware device connections: device_id -> set of WebSocket connections
_hardware_clients: Dict[str, Set[WebSocket]] = {}

# Track browser connections (for payment notifications in UI)
_browser_clients: Dict[str, Set[WebSocket]] = {}


def get_connected_device_ids() -> list[str]:
"""Return list of device IDs with active WebSocket connections."""
return list(_connected_clients.keys())
"""Return list of device IDs with active hardware connections."""
return list(_hardware_clients.keys())


def is_device_connected(device_id: str) -> bool:
"""Check if a device has any active WebSocket connections."""
return device_id in _connected_clients and len(_connected_clients[device_id]) > 0
"""Check if a hardware device has any active WebSocket connections."""
return device_id in _hardware_clients and len(_hardware_clients[device_id]) > 0


async def send_to_device(device_id: str, message: str) -> bool:
"""
Send a message to all WebSocket connections for a device.
Send a message to all WebSocket connections for a device (hardware + browser).
Returns True if message was sent to at least one client.
"""
if device_id not in _connected_clients:
logger.warning(f"No WebSocket connections for device {device_id}")
return False

sent = False
dead_connections: Set[WebSocket] = set()

for websocket in _connected_clients[device_id]:
try:
await websocket.send_text(message)
sent = True
except Exception as e:
logger.debug(f"Failed to send to WebSocket: {e}")
dead_connections.add(websocket)

# Clean up dead connections
for ws in dead_connections:
_connected_clients[device_id].discard(ws)

# Remove device entry if no connections left
if device_id in _connected_clients and not _connected_clients[device_id]:
del _connected_clients[device_id]
# Send to hardware clients
if device_id in _hardware_clients:
dead_connections: Set[WebSocket] = set()
for websocket in _hardware_clients[device_id]:
try:
await websocket.send_text(message)
sent = True
except Exception as e:
logger.debug(f"Failed to send to hardware: {e}")
dead_connections.add(websocket)
for ws in dead_connections:
_hardware_clients[device_id].discard(ws)
if not _hardware_clients[device_id]:
del _hardware_clients[device_id]

# Send to browser clients (so UI shows payment received)
if device_id in _browser_clients:
dead_connections = set()
for websocket in _browser_clients[device_id]:
try:
await websocket.send_text(message)
sent = True
except Exception as e:
logger.debug(f"Failed to send to browser: {e}")
dead_connections.add(websocket)
for ws in dead_connections:
_browser_clients[device_id].discard(ws)
if not _browser_clients[device_id]:
del _browser_clients[device_id]

if not sent:
logger.warning(f"No WebSocket connections for device {device_id}")

return sent


@devicetimer_websocket_router.websocket("/api/v1/ws/{device_id}")
async def websocket_endpoint(websocket: WebSocket, device_id: str):
async def websocket_endpoint(
websocket: WebSocket,
device_id: str,
type: Optional[str] = Query(default="hardware")
):
"""
WebSocket endpoint for device connections.
Hardware devices connect here to receive payment notifications.

Query params:
type: "hardware" (default) for ESP32 devices, "browser" for UI connections
"""
await websocket.accept()

# Select the appropriate client pool
is_browser = type == "browser"
clients = _browser_clients if is_browser else _hardware_clients
client_type = "browser" if is_browser else "hardware"

# Add to tracking
if device_id not in _connected_clients:
_connected_clients[device_id] = set()
_connected_clients[device_id].add(websocket)
if device_id not in clients:
clients[device_id] = set()
clients[device_id].add(websocket)

logger.info(f"Device {device_id} connected. Total connections: {len(_connected_clients[device_id])}")
logger.info(f"{client_type.capitalize()} connected for device {device_id}")

try:
while True:
# Keep connection alive, wait for messages (ping/pong handled automatically)
data = await websocket.receive_text()
# Hardware might send status updates, we just acknowledge
logger.debug(f"Received from {device_id}: {data}")
logger.debug(f"Received from {device_id} ({client_type}): {data}")
except WebSocketDisconnect:
logger.info(f"Device {device_id} disconnected")
logger.info(f"{client_type.capitalize()} disconnected for device {device_id}")
except Exception as e:
logger.debug(f"WebSocket error for {device_id}: {e}")
finally:
# Remove from tracking
if device_id in _connected_clients:
_connected_clients[device_id].discard(websocket)
if not _connected_clients[device_id]:
del _connected_clients[device_id]
logger.info(f"Device {device_id} cleaned up. Connected devices: {list(_connected_clients.keys())}")
if device_id in clients:
clients[device_id].discard(websocket)
if not clients[device_id]:
del clients[device_id]
logger.debug(f"Connected hardware devices: {list(_hardware_clients.keys())}")


@devicetimer_websocket_router.get("/api/v1/ws/status")
Expand Down