/
/
/
1"""WebSocket client handler for Music Assistant API."""
2
3from __future__ import annotations
4
5import asyncio
6import contextvars
7import inspect
8import logging
9from concurrent import futures
10from contextlib import suppress
11from functools import partial
12from typing import TYPE_CHECKING, Any, Final
13from uuid import uuid4
14
15from aiohttp import WSMsgType, web
16from music_assistant_models.api import (
17 CommandMessage,
18 ErrorResultMessage,
19 MessageType,
20 SuccessResultMessage,
21)
22from music_assistant_models.auth import AuthProviderType, Scope, User
23from music_assistant_models.enums import EventType
24from music_assistant_models.errors import (
25 AuthenticationRequired,
26 InsufficientPermissions,
27 InvalidCommand,
28 InvalidToken,
29 MusicAssistantError,
30)
31from music_assistant_models.event import MassEvent
32from music_assistant_models.media_items.metadata import IMAGE_PROXY_ID_RESOLVER
33from music_assistant_models.translations import TRANSLATION_RESOLVER
34
35from music_assistant.constants import HOMEASSISTANT_SYSTEM_USER, VERBOSE_LOG_LEVEL
36from music_assistant.helpers.api import APICommandHandler, parse_arguments
37
38from .helpers.auth_middleware import (
39 has_scope,
40 is_request_from_ingress,
41 resolve_command_impersonation,
42 set_current_client_id,
43 set_current_token,
44 set_current_user,
45 set_impersonated_user,
46 set_sendspin_player_id,
47)
48from .helpers.auth_providers import get_ha_user_details, get_ha_user_role
49
50if TYPE_CHECKING:
51 from music_assistant.controllers.webserver import WebserverController
52
53MAX_PENDING_MSG = 512
54CANCELLATION_ERRORS: Final = (asyncio.CancelledError, futures.CancelledError)
55
56
57class WebsocketClientHandler:
58 """Handle an active websocket client connection."""
59
60 def __init__(self, webserver: WebserverController, request: web.Request) -> None:
61 """Initialize an active connection."""
62 self.webserver = webserver
63 self.mass = webserver.mass
64 self.request = request
65 self.client_id = uuid4().hex
66 self.wsock = web.WebSocketResponse(heartbeat=25)
67 self._to_write: asyncio.Queue[str | None] = asyncio.Queue(maxsize=MAX_PENDING_MSG)
68 self._handle_task: asyncio.Task[Any] | None = None
69 self._writer_task: asyncio.Task[None] | None = None
70 self._logger = webserver.logger
71 self._authenticated_user: User | None = (
72 None # Will be set after auth command or from Ingress
73 )
74 self._current_token: str | None = None # Will be set after auth command
75 self._token_id: str | None = None # Will be set after auth for tracking revocation
76 self._sendspin_player_id: str | None = None # Set if client is a sendspin web player
77 self._locale: str | None = None # UI locale declared by the client (auth arg / set_locale)
78 self._is_ingress = is_request_from_ingress(request)
79 self._events_unsub_callback: Any = None # Will be set after authentication
80 # Track WebRTC session ID if this is a WebRTC gateway connection
81 self._webrtc_session_id: str | None = request.query.get("webrtc_session_id")
82 # try to dynamically detect the base_url of a client if proxied or behind Ingress
83 self.base_url: str | None = None
84 if forward_host := request.headers.get("X-Forwarded-Host"):
85 ingress_path = request.headers.get("X-Ingress-Path", "")
86 forward_proto = request.headers.get("X-Forwarded-Proto", request.protocol)
87 self.base_url = f"{forward_proto}://{forward_host}{ingress_path}"
88
89 async def disconnect(self) -> None:
90 """Disconnect client."""
91 self._cancel()
92 if self._writer_task is not None:
93 await self._writer_task
94
95 async def handle_client(self) -> web.WebSocketResponse:
96 """Handle a websocket response."""
97 # ruff: noqa: PLR0915
98 request = self.request
99 wsock = self.wsock
100 try:
101 async with asyncio.timeout(10):
102 await wsock.prepare(request)
103 except TimeoutError:
104 self._logger.warning("Timeout preparing request from %s", request.remote)
105 return wsock
106
107 self._logger.log(VERBOSE_LOG_LEVEL, "Connection from %s", request.remote)
108 self._handle_task = asyncio.current_task()
109 self._writer_task = self.mass.create_task(self._writer())
110
111 # send server(version) info when client connects
112 server_info = self.mass.get_server_info()
113 await self._send_message(server_info)
114
115 # Block until onboarding is complete
116 if not self.webserver.auth.has_users and not self._is_ingress:
117 await self._send_message(
118 ErrorResultMessage(
119 "connection", 503, "Setup required", translation_key="setup_required"
120 )
121 )
122 await wsock.close()
123 return wsock
124
125 # For Ingress connections, auto-create/link user and subscribe to events immediately
126 # For regular connections, events will be subscribed after successful authentication
127 if self._is_ingress:
128 await self._handle_ingress_auth()
129 self._subscribe_to_events()
130
131 disconnect_warn = None
132
133 try:
134 while not wsock.closed:
135 msg = await wsock.receive()
136
137 if msg.type in (WSMsgType.CLOSE, WSMsgType.CLOSING, WSMsgType.CLOSED):
138 break
139
140 if msg.type != WSMsgType.TEXT:
141 continue
142
143 self._logger.log(VERBOSE_LOG_LEVEL, "Received: %s", msg.data)
144
145 try:
146 command_msg = CommandMessage.from_json(msg.data)
147 except ValueError:
148 disconnect_warn = f"Received invalid JSON: {msg.data}"
149 break
150
151 await self._handle_command(command_msg)
152
153 except asyncio.CancelledError:
154 self._logger.debug("Connection closed by client")
155
156 except Exception:
157 self._logger.exception("Unexpected error inside websocket API")
158
159 finally:
160 # Handle connection shutting down.
161 if self._events_unsub_callback:
162 self._events_unsub_callback()
163 self._logger.log(VERBOSE_LOG_LEVEL, "Unsubscribed from events")
164
165 # Unregister from webserver tracking
166 self.webserver.unregister_websocket_client(self)
167
168 # Drop any dashboard registrations owned by this connection
169 self.mass.dashboard.handle_client_disconnected(self.client_id)
170
171 try:
172 self._to_write.put_nowait(None)
173 # Make sure all error messages are written before closing
174 await self._writer_task
175 await wsock.close()
176 except asyncio.QueueFull: # can be raised by put_nowait
177 self._writer_task.cancel()
178
179 finally:
180 if disconnect_warn is None:
181 self._logger.log(VERBOSE_LOG_LEVEL, "Disconnected")
182 else:
183 self._logger.warning("Disconnected: %s", disconnect_warn)
184
185 return wsock
186
187 async def _handle_command(self, msg: CommandMessage) -> None:
188 """Handle an incoming command from the client."""
189 self._logger.log(VERBOSE_LOG_LEVEL, "Handling command %s", msg.command)
190
191 # Handle special "auth" command
192 if msg.command == "auth":
193 await self._handle_auth_command(msg)
194 return
195
196 # Handle special "translations/set_locale" command (updates connection state)
197 if msg.command == "translations/set_locale":
198 await self._handle_set_locale_command(msg)
199 return
200
201 # work out handler for the given path/command
202 handler = self.mass.command_handlers.get(msg.command)
203
204 if handler is None:
205 await self._send_message(
206 ErrorResultMessage(
207 msg.message_id,
208 InvalidCommand.error_code,
209 f"Invalid command: {msg.command}",
210 translation_key="invalid_command",
211 )
212 )
213 self._logger.warning("Invalid command: %s", msg.command)
214 return
215
216 # Put this connection's identity in context for the API methods. ContextVars live
217 # for as long as the connection does, so every command sets all of them: an
218 # unauthenticated handler must see this connection's own (possibly absent) user
219 # rather than whatever the command before it left behind.
220 set_current_client_id(self.client_id)
221 set_current_user(self._authenticated_user)
222 set_current_token(self._current_token)
223 set_sendspin_player_id(self._sendspin_player_id)
224
225 # Check authentication if required
226 if handler.authenticated or handler.required_scope:
227 # For Ingress, user should already be set from _handle_ingress_auth
228 # For regular connections, user must be set via auth command
229 if self._authenticated_user is None:
230 await self._send_message(
231 ErrorResultMessage(
232 msg.message_id,
233 AuthenticationRequired.error_code,
234 "Authentication required. Please send auth command first.",
235 translation_key="authentication_required",
236 )
237 )
238 return
239
240 # Check scope if required
241 if handler.required_scope and not has_scope(
242 self._authenticated_user, handler.required_scope
243 ):
244 await self._send_message(
245 ErrorResultMessage(
246 msg.message_id,
247 InsufficientPermissions.error_code,
248 f"This command requires the {handler.required_scope} scope",
249 translation_key="insufficient_permissions",
250 )
251 )
252 return
253
254 # schedule task to handle the command
255 self.mass.create_task(self._run_handler(handler, msg))
256
257 async def _run_handler(self, handler: APICommandHandler, msg: CommandMessage) -> None:
258 """Run command handler and send response."""
259 try:
260 # handle the optional impersonation argument for impersonation-enabled commands
261 if handler.allow_impersonation and msg.args:
262 if impersonation_user := await resolve_command_impersonation(self.mass, msg.args):
263 set_impersonated_user(impersonation_user)
264 args = parse_arguments(handler.signature, handler.type_hints, msg.args)
265 result: Any = handler.target(**args)
266 if hasattr(result, "__anext__"):
267 # handle async generator (for really large listings)
268 items: list[Any] = []
269 async for item in result:
270 items.append(item)
271 if len(items) >= 500:
272 await self._send_message(
273 SuccessResultMessage(msg.message_id, items, partial=True)
274 )
275 items = []
276 result = items
277 elif inspect.iscoroutine(result):
278 result = await result
279 await self._send_message(SuccessResultMessage(msg.message_id, result))
280 except MusicAssistantError as err:
281 # Expected operational errors (player unavailable, queue empty, etc.)
282 # Log at warning level since these are normal error responses, not crashes.
283 self._logger.warning("%s: %s", msg.command, err)
284 err_msg = str(err) or err.__class__.__name__
285 # err_msg is the English fallback; the translation_key (per-type default or a
286 # provider override) localizes `details` to the connection locale at serialization.
287 await self._send_message(
288 ErrorResultMessage(
289 msg.message_id,
290 err.error_code,
291 err_msg,
292 translation_key=err.translation_key,
293 translation_args=err.translation_args,
294 translation_owner=err.translation_owner,
295 )
296 )
297 except Exception as err:
298 if self._logger.isEnabledFor(logging.DEBUG):
299 self._logger.exception("Error handling message: %s", msg)
300 else:
301 self._logger.error("Error handling message: %s: %s", msg.command, str(err))
302 err_msg = str(err) or err.__class__.__name__
303 await self._send_message(
304 ErrorResultMessage(msg.message_id, getattr(err, "error_code", 999), err_msg)
305 )
306
307 async def _writer(self) -> None:
308 """Write outgoing messages."""
309 # Exceptions if Socket disconnected or cancelled by connection handler
310 with suppress(RuntimeError, ConnectionResetError, *CANCELLATION_ERRORS):
311 while not self.wsock.closed:
312 if (process := await self._to_write.get()) is None:
313 break
314
315 if callable(process):
316 message: str = process()
317 else:
318 message = process
319 self._logger.log(VERBOSE_LOG_LEVEL, "Writing: %s", message)
320 await self.wsock.send_str(message)
321
322 async def _send_message(self, message: MessageType) -> None:
323 """
324 Send a message to the client (for large response messages).
325
326 Runs JSON serialization in executor to avoid blocking for large messages.
327 Closes connection if the client is not reading the messages.
328
329 Async friendly.
330 """
331 # Run JSON serialization in executor to avoid blocking for large messages.
332 # copy_context() propagates the IMAGE_PROXY_ID_RESOLVER and TRANSLATION_RESOLVER
333 # ContextVars into the executor thread so that nested models can inject `proxy_id`
334 # and localize human-readable fields via their `__post_serialize__` hooks.
335 loop = asyncio.get_running_loop()
336 token = IMAGE_PROXY_ID_RESOLVER.set(self.mass.metadata.compute_image_id)
337 token_loc = TRANSLATION_RESOLVER.set(
338 partial(self.mass.translations.get_translation, locale=self._locale)
339 )
340 try:
341 ctx = contextvars.copy_context()
342 _message = await loop.run_in_executor(None, ctx.run, message.to_json)
343 finally:
344 IMAGE_PROXY_ID_RESOLVER.reset(token)
345 TRANSLATION_RESOLVER.reset(token_loc)
346
347 try:
348 self._to_write.put_nowait(_message)
349 except asyncio.QueueFull:
350 self._logger.error("Client exceeded max pending messages: %s", MAX_PENDING_MSG)
351
352 self._cancel()
353
354 def _send_message_sync(self, message: MessageType) -> None:
355 """
356 Send a message from a sync context (for small messages like events).
357
358 Serializes inline without executor overhead since events are typically small.
359 """
360 token = IMAGE_PROXY_ID_RESOLVER.set(self.mass.metadata.compute_image_id)
361 token_loc = TRANSLATION_RESOLVER.set(
362 partial(self.mass.translations.get_translation, locale=self._locale)
363 )
364 try:
365 _message = message.to_json()
366 finally:
367 IMAGE_PROXY_ID_RESOLVER.reset(token)
368 TRANSLATION_RESOLVER.reset(token_loc)
369
370 try:
371 self._to_write.put_nowait(_message)
372 except asyncio.QueueFull:
373 self._logger.error("Client exceeded max pending messages: %s", MAX_PENDING_MSG)
374
375 self._cancel()
376
377 async def _handle_auth_command(self, msg: CommandMessage) -> None:
378 """
379 Handle WebSocket authentication command.
380
381 :param msg: The auth command message with access token.
382 """
383 # Extract token from args (support both 'token' and 'access_token' for backward compat)
384 token = msg.args.get("token") if msg.args else None
385 if not token:
386 token = msg.args.get("access_token") if msg.args else None
387 if not token:
388 await self._send_message(
389 ErrorResultMessage(
390 msg.message_id,
391 AuthenticationRequired.error_code,
392 "token required in args",
393 )
394 )
395 return
396
397 # Authenticate with token
398 user = await self.webserver.auth.authenticate_with_token(token)
399 if not user:
400 await self._send_message(
401 ErrorResultMessage(
402 msg.message_id,
403 InvalidToken.error_code,
404 "Invalid or expired token",
405 translation_key="invalid_token",
406 )
407 )
408 return
409
410 # Security: Deny homeassistant system user on regular (non-Ingress) webserver
411 if not self._is_ingress and user.username == HOMEASSISTANT_SYSTEM_USER:
412 await self._send_message(
413 ErrorResultMessage(
414 msg.message_id,
415 InvalidToken.error_code,
416 "Home Assistant system user not allowed on regular webserver",
417 )
418 )
419 return
420
421 # Get token_id for tracking revocation events
422 token_id = await self.webserver.auth.get_token_id_from_token(token)
423
424 # Store authenticated user, token, and token_id
425 self._authenticated_user = user
426 self._current_token = token
427 self._token_id = token_id
428 self._logger.info("WebSocket client authenticated as %s", user.username)
429
430 # Optionally store the UI locale declared with the auth command and warm it up
431 if msg.args and (locale := msg.args.get("locale")):
432 self._locale = locale
433 await self.mass.translations.ensure_locale_loaded(locale)
434
435 # Send success response
436 await self._send_message(
437 SuccessResultMessage(
438 msg.message_id,
439 {"authenticated": True, "user": user.to_dict()},
440 )
441 )
442
443 # Subscribe to events after successful authentication
444 self._subscribe_to_events()
445
446 # Register with webserver for tracking
447 self.webserver.register_websocket_client(self)
448
449 async def _handle_set_locale_command(self, msg: CommandMessage) -> None:
450 """
451 Handle the WebSocket set_locale command (updates the connection's UI locale).
452
453 :param msg: The set_locale command message; expects a "locale" arg.
454 """
455 locale = msg.args.get("locale") if msg.args else None
456 if not locale:
457 await self._send_message(
458 ErrorResultMessage(
459 msg.message_id,
460 InvalidCommand.error_code,
461 "locale required in args",
462 )
463 )
464 return
465 self._locale = locale
466 await self.mass.translations.ensure_locale_loaded(locale)
467 await self._send_message(SuccessResultMessage(msg.message_id, {"locale": locale}))
468
469 async def _handle_ingress_auth(self) -> None:
470 """Handle authentication for Ingress connections (auto-create/link user)."""
471 ingress_user_id = self.request.headers.get("X-Remote-User-ID")
472 ingress_username = self.request.headers.get("X-Remote-User-Name")
473 ingress_display_name = self.request.headers.get("X-Remote-User-Display-Name")
474
475 if ingress_user_id and ingress_username:
476 # Try to find existing user linked to this HA user ID
477 user = await self.webserver.auth.get_user_by_provider_link(
478 AuthProviderType.HOME_ASSISTANT, ingress_user_id
479 )
480
481 if not user:
482 # Check if a user with this username already exists
483 user = await self.webserver.auth.get_user_by_username(ingress_username)
484
485 if not user:
486 # New user - fetch details from HA
487 ha_username, ha_display_name, avatar_url = await get_ha_user_details(
488 self.mass, ingress_user_id
489 )
490 # Auto-create user for Ingress (they're already authenticated by HA)
491 role = await get_ha_user_role(self.mass, ingress_user_id)
492 user = await self.webserver.auth.create_user(
493 username=ha_username or ingress_username,
494 role=role,
495 display_name=ha_display_name or ingress_display_name,
496 avatar_url=avatar_url,
497 )
498
499 # Link to Home Assistant provider (or create the link if user already existed)
500 await self.webserver.auth.link_user_to_provider(
501 user, AuthProviderType.HOME_ASSISTANT, ingress_user_id
502 )
503
504 # Update user with HA details if available (HA is source of truth)
505 # Fall back to ingress headers if API lookup doesn't return values
506 _, ha_display_name, avatar_url = await get_ha_user_details(self.mass, ingress_user_id)
507 final_display_name = ha_display_name or ingress_display_name
508 if final_display_name or avatar_url:
509 user = await self.webserver.auth.update_user(
510 user,
511 display_name=final_display_name,
512 avatar_url=avatar_url,
513 )
514
515 self._authenticated_user = user
516 self._logger.debug("Ingress user authenticated: %s", user.username)
517 else:
518 # No HA user headers - allow homeassistant system user to connect with token
519 # This allows the Home Assistant integration to connect via the internal network
520 # The token authentication happens in _handle_auth_message
521 self._logger.debug("Ingress connection without user headers, expecting token auth")
522
523 def _subscribe_to_events(self) -> None:
524 """Subscribe to Mass events and forward them to the client."""
525 if self._events_unsub_callback is not None:
526 # Already subscribed
527 return
528
529 def handle_event(event: MassEvent) -> None:
530 # filter events for objects the user has no access to
531 if (
532 self._authenticated_user
533 and self._authenticated_user.player_filter
534 and event.event
535 in (
536 EventType.PLAYER_ADDED,
537 EventType.PLAYER_REMOVED,
538 EventType.PLAYER_UPDATED,
539 EventType.PLAYER_SLEEP_TIMER_UPDATED,
540 EventType.QUEUE_ADDED,
541 EventType.QUEUE_ITEMS_UPDATED,
542 EventType.QUEUE_TIME_UPDATED,
543 EventType.QUEUE_UPDATED,
544 )
545 and event.object_id
546 and event.object_id not in self._authenticated_user.player_filter
547 and event.object_id != self._sendspin_player_id
548 ):
549 return
550
551 if event.event == EventType.SETUP_FLOW_UPDATED:
552 # setup flow steps carry prefilled values, OAuth urls and the
553 # flow_id guarding the unauthenticated callback route - only
554 # users who could interact with the flow may receive them
555 user = self._authenticated_user
556 if user is None:
557 return
558 required = (
559 self.mass.config.get_setup_flow_required_scope(event.object_id)
560 if event.object_id
561 else None
562 )
563 if required is None:
564 # flow already popped (terminal step race): the flow kind is no
565 # longer known, so require both config scopes to be safe
566 if not has_scope(user, Scope.CONFIG_PROVIDERS_WRITE) or not has_scope(
567 user, Scope.CONFIG_PLAYERS_WRITE
568 ):
569 return
570 elif not has_scope(user, required):
571 return
572
573 if event.event == EventType.TASKS_UPDATED:
574 if self._authenticated_user is None:
575 return
576 task_data = self.mass.tasks.list_tasks_for_user(self._authenticated_user)
577 self._send_message_sync(
578 MassEvent(
579 event=event.event,
580 object_id=event.object_id,
581 data=task_data,
582 )
583 )
584 return
585
586 self._send_message_sync(event)
587
588 self._events_unsub_callback = self.mass.subscribe(handle_event)
589 self._logger.debug("Subscribed to events")
590
591 def _cancel(self) -> None:
592 """Cancel the connection."""
593 if self._handle_task is not None:
594 self._handle_task.cancel()
595 if self._writer_task is not None:
596 self._writer_task.cancel()
597