GlcoinCore/blitz_api
1
0
Fork 0
mirror of https://github.com/fusion44/blitz_api.git synced 2026-08-13 11:52:45 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 1

feat(api): serve realtime updates over /ws WebSocket

Browse source 
Replace the /sse/subscribe endpoint with a /ws WebSocket endpoint backed
by ws_mgr, and port the warmup helper (_send_sse_event -> _send_ws_event,
SSE -> Event, broadcast_sse_msg -> broadcast_msg). Also fixes two tests
left over from the prior SSE->Event/broadcast_msg rename that still
referenced the old names and were failing collection/execution.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
fusion44 2026-07-11 12:07:21 +02:00
parent 32414e6932
commit 5089925c88
No known key found for this signature in database
4 changed files with 79 additions and 62 deletions
Download patch file Download diff file Expand all files Collapse all files

100
app/main.py
View file

@ -3,7 +3,7 @@ import sys
import traceback
from contextlib import asynccontextmanager
from fastapi import FastAPI, Request
from fastapi import FastAPI, Request, WebSocket
from fastapi.encoders import jsonable_encoder
from fastapi.exceptions import HTTPException, RequestValidationError
from loguru import logger
@ -12,19 +12,20 @@ from redis.asyncio import Redis
from starlette import status
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import JSONResponse, RedirectResponse
from starlette.websockets import WebSocketDisconnect
from app.api.config import config as dconfig
from app.api.error_report.report import Report
from app.api.models import ApiErrors, ApiStartupStatus, ErrorMessage, StartupState
from app.api.utils import SSE, broadcast_sse_msg, build_sse_event, sse_mgr
from app.api.utils import Event, broadcast_msg
from app.api.warmup import (
get_bitcoin_client_warmup_data,
get_full_client_warmup_data,
get_full_client_warmup_data_bitcoinonly,
)
from app.api.ws_manager import ws_mgr
from app.apps.router import register_app_status_update_handlers
from app.apps.router import router as app_router
from app.auth.auth_bearer import JWTBearer
from app.auth.auth_handler import (
handle_local_cookie,
register_cookie_updater,
@ -97,7 +98,7 @@ async def lifespan(app: FastAPI):
raise RuntimeError("Redis not initialized correctly, got a Sentinel")
register_cookie_updater()
await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump())
await broadcast_msg(Event.SYSTEM_STARTUP_INFO, api_startup_status.model_dump())
btc_task = asyncio.create_task(_initialize_bitcoin())
ln_task = asyncio.create_task(_initialize_lightning())
await register_all_handlers()
@ -191,7 +192,7 @@ async def _set_startup_status(
api_startup_status.lightning_msg = lightning_msg
asyncio.create_task(warmup_new_connections())
await broadcast_sse_msg(SSE.SYSTEM_STARTUP_INFO, api_startup_status.model_dump())
await broadcast_msg(Event.SYSTEM_STARTUP_INFO, api_startup_status.model_dump())
@logger.catch
@ -278,45 +279,32 @@ def index(req: Request):
new_connections = []
async def _send_sse_event(id, event, data):
return await sse_mgr.send_to_single(id, build_sse_event(event, data))
async def _send_ws_event(id, event, data):
return await ws_mgr.send_to_single(id, event, data)
@app.get(
"/sse/subscribe",
status_code=status.HTTP_200_OK,
)
async def stream(request: Request):
token = request.cookies.get("access_token")
if not token:
# No token in cookies found, try to get it from the Authorization header
token = request.headers.get("authorization")
@app.websocket("/ws")
async def stream(websocket: WebSocket):
conn_id, authed = await ws_mgr.connect(websocket)
if not authed:
return # ws_mgr already closed the socket
if not token:
# Raise an exception as there is not token found
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="No authorization code.",
)
token = token.replace("Bearer ", "")
if not JWTBearer().verify_jwt(jwtoken=token):
# token is invalid
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="Invalid authorization code.",
)
event_source, id = sse_mgr.add_connection(request)
new_connections.append(id)
await _send_sse_event(
id, SSE.SYSTEM_STARTUP_INFO, jsonable_encoder(api_startup_status.model_dump())
new_connections.append(conn_id)
await _send_ws_event(
conn_id,
Event.SYSTEM_STARTUP_INFO,
jsonable_encoder(api_startup_status.model_dump()),
)
asyncio.create_task(warmup_new_connections())
return event_source
try:
while True:
# server-push only; receive loop just detects disconnect
await websocket.receive_text()
except WebSocketDisconnect:
ws_mgr.disconnect(conn_id)
if conn_id in new_connections:
new_connections.remove(conn_id)
warmup_running = False
@ -331,11 +319,11 @@ async def warmup_new_connections():
async def _handle(id, event, res):
match res:
case BaseModel():
return await _send_sse_event(id, event, res.model_dump())
return await _send_ws_event(id, event, res.model_dump())
case dict() | list():
return await _send_sse_event(id, event, res)
return await _send_ws_event(id, event, res)
case Ok(data) if data and isinstance(data, BaseModel):
return await _send_sse_event(id, event, data.model_dump())
return await _send_ws_event(id, event, data.model_dump())
case Ok(data) if not data:
logger.debug(f"No data to send for warmup event {event}")
return
@ -346,7 +334,7 @@ async def warmup_new_connections():
)
logger.error(f"Error while fetching warmup_data for {event}: {res}")
return await _send_sse_event(id, event, {"error": f"{res}"})
return await _send_ws_event(id, event, {"error": f"{res}"})
global new_connections
if len(new_connections) == 0:
@ -368,13 +356,13 @@ async def warmup_new_connections():
for id in new_connections:
await asyncio.gather(
*[
_handle(id, SSE.SYSTEM_INFO, res[0]),
_handle(id, SSE.BTC_INFO, res[1]),
_handle(id, SSE.LN_INFO, res[2]),
_handle(id, SSE.LN_FEE_REVENUE, res[3]),
_handle(id, SSE.WALLET_BALANCE, res[4]),
_handle(id, SSE.APP_STATE_MESSAGE, res[5]),
_handle(id, SSE.HARDWARE_INFO, res[6]),
_handle(id, Event.SYSTEM_INFO, res[0]),
_handle(id, Event.BTC_INFO, res[1]),
_handle(id, Event.LN_INFO, res[2]),
_handle(id, Event.LN_FEE_REVENUE, res[3]),
_handle(id, Event.WALLET_BALANCE, res[4]),
_handle(id, Event.APP_STATE_MESSAGE, res[5]),
_handle(id, Event.HARDWARE_INFO, res[6]),
]
)
@ -384,10 +372,10 @@ async def warmup_new_connections():
for id in new_connections:
await asyncio.gather(
*[
_handle(id, SSE.SYSTEM_INFO, res[0]),
_handle(id, SSE.BTC_INFO, res[1]),
_handle(id, SSE.APP_STATE_MESSAGE, res[2]),
_handle(id, SSE.HARDWARE_INFO, res[3]),
_handle(id, Event.SYSTEM_INFO, res[0]),
_handle(id, Event.BTC_INFO, res[1]),
_handle(id, Event.APP_STATE_MESSAGE, res[2]),
_handle(id, Event.HARDWARE_INFO, res[3]),
]
)
@ -402,8 +390,8 @@ async def warmup_new_connections():
for id in new_connections:
await asyncio.gather(
*[
_handle(id, SSE.BTC_INFO, res[0]),
_handle(id, SSE.HARDWARE_INFO, res[1]),
_handle(id, Event.BTC_INFO, res[0]),
_handle(id, Event.HARDWARE_INFO, res[1]),
]
)
@ -418,7 +406,7 @@ async def warmup_new_connections():
# Bitcoin Core and Lightning running
res = await get_hardware_info()
for id in new_connections:
await _send_sse_event(id, SSE.HARDWARE_INFO, res)
await _send_ws_event(id, Event.HARDWARE_INFO, res)
# don't clear new_connections,
# we'll try again later when api is initialized

2
tests/test_hw_gatherer_resilience.py
View file

@ -33,7 +33,7 @@ async def test_gatherer_survives_get_hardware_info_error(monkeypatch):
raise asyncio.CancelledError()
monkeypatch.setattr(service, "get_hardware_info", flaky_hardware_info)
monkeypatch.setattr(service, "broadcast_sse_msg", fake_broadcast)
monkeypatch.setattr(service, "broadcast_msg", fake_broadcast)
monkeypatch.setattr(service.asyncio, "sleep", fake_sleep)
with pytest.raises(asyncio.CancelledError):

10
tests/test_warmup_events.py
View file

@ -7,7 +7,7 @@ name the WebUI does not listen to.
"""
from app.api.models import ApiStartupStatus, StartupState
from app.api.utils import SSE
from app.api.utils import Event
from app.apps.constants import AppManagementProcessState
from app.apps.models import AppStatusQueryResult, AppStatusUpdateTaskMessage
from app.external.result_type.src.result.result import Ok
@ -37,7 +37,7 @@ def _patch_common(monkeypatch, sent, *, node_type: str, lightning: StartupState)
"api_startup_status",
ApiStartupStatus(bitcoin=StartupState.DONE, lightning=lightning),
)
monkeypatch.setattr(main, "_send_sse_event", fake_send)
monkeypatch.setattr(main, "_send_ws_event", fake_send)
monkeypatch.setattr(
main, "get_bitcoin_client_warmup_data", fake_bitcoin_warmup_data
)
@ -46,7 +46,7 @@ def _patch_common(monkeypatch, sent, *, node_type: str, lightning: StartupState)
async def test_bitcoinonly_warmup_sends_app_status_as_app_state_message(monkeypatch):
"""The WebUI only listens for SSE.APP_STATE_MESSAGE to populate the Apps
"""The WebUI only listens for Event.APP_STATE_MESSAGE to populate the Apps
tab. In bitcoin-only mode the warmup app status must be sent under that
event, exactly like in lightning mode."""
sent = []
@ -69,8 +69,8 @@ async def test_bitcoinonly_warmup_sends_app_status_as_app_state_message(monkeypa
await main.warmup_new_connections()
events = [event for (_, event, _) in sent]
assert SSE.APP_STATE_MESSAGE in events, (
f"expected app status under event '{SSE.APP_STATE_MESSAGE}', "
assert Event.APP_STATE_MESSAGE in events, (
f"expected app status under event '{Event.APP_STATE_MESSAGE}', "
f"got events: {events}"
)

29
tests/test_ws_endpoint.py Normal file
View file

@ -0,0 +1,29 @@
import json
from starlette.testclient import TestClient
from starlette.websockets import WebSocketDisconnect
from app.auth.auth_handler import sign_jwt
from app.main import app
client = TestClient(app)
def test_ws_valid_auth_receives_startup_frame():
# conftest.py sets BAPI_JWT_SECRET/ALGORITHM, so sign_jwt() is a valid token
token = sign_jwt()
with client.websocket_connect("/ws") as ws:
ws.send_text(json.dumps({"type": "auth", "token": token}))
frame = json.loads(ws.receive_text())
assert frame["event"] == "system_startup_info"
assert "data" in frame
def test_ws_bad_token_is_closed_4401():
with client.websocket_connect("/ws") as ws:
ws.send_text(json.dumps({"type": "auth", "token": "not-a-jwt"}))
try:
ws.receive_text()
assert False, "expected the server to close the connection"
except WebSocketDisconnect as e:
assert e.code == 4401