wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Port WebSocket proxies to the new BaseTask API

Browse Source 
`WebSocketProxyServerAgent` / `WebSocketProxyClientAgent` are
renamed to `WebSocketProxyServerTask` / `WebSocketProxyClientTask`
and updated for the post-refactor surface:

- Drop `bus=` from the constructor; the bus arrives via
  `BaseTask.attach` from the runner.
- Constructor params `agent_name` / `remote_agent_name` /
  `local_agent_name` → `task_name` / `remote_task_name` /
  `local_task_name` (matching `BusBridgeProcessor`).
- Move setup logic from the now-removed `on_ready` hook into
  `start()`; replace `_stop()` overrides with `stop()`.
- Add `_handle_task_end` / `_handle_task_cancel` overrides that
  set `_finished_event` so `PipelineRunner._cancel_spawned_tasks`
  can drive these bus-only tasks to a clean exit.
- Update the registry-message field reference
  (`agents=`/`message.agents` → `tasks=`/`message.tasks`)
  and `TaskReadyData.task_name` access.
- Tighten the server's `_send_ws` exception handling to only
  catch `WebSocketDisconnect`.
- Update install hints (`pipecat-ai[websockets-base]` for the
  client, `starlette` for the server) and refresh docstrings/
  examples to use `runner.spawn(...)`.
This commit is contained in:
Aleix Conchillo Flaqué
2026-05-13 22:24:37 -07:00
parent 86f8f137a8
commit 5f86e39038
4 changed files with 159 additions and 144 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/pipecat/tasks/proxy/__init__.py
View File

@@ -4,14 +4,14 @@
# SPDX-License-Identifier: BSD 2-Clause License # SPDX-License-Identifier: BSD 2-Clause License
# #
"""Proxy agents for forwarding bus messages over network transports.""" """Proxy tasks for forwarding bus messages over network transports."""
from pipecat.tasks.proxy.websocket import ( from pipecat.tasks.proxy.websocket import (
WebSocketProxyClientAgent, WebSocketProxyClientTask,
WebSocketProxyServerAgent, WebSocketProxyServerTask,
) )
__all__ = [ __all__ = [
"WebSocketProxyClientAgent", "WebSocketProxyClientTask",
"WebSocketProxyServerAgent", "WebSocketProxyServerTask",
] ]

10
src/pipecat/tasks/proxy/websocket/__init__.py
View File

@@ -4,12 +4,12 @@
# SPDX-License-Identifier: BSD 2-Clause License # SPDX-License-Identifier: BSD 2-Clause License
# #
"""WebSocket proxy agents for forwarding bus messages.""" """WebSocket proxy tasks for forwarding bus messages."""
from pipecat.tasks.proxy.websocket.client import WebSocketProxyClientAgent from pipecat.tasks.proxy.websocket.client import WebSocketProxyClientTask
from pipecat.tasks.proxy.websocket.server import WebSocketProxyServerAgent from pipecat.tasks.proxy.websocket.server import WebSocketProxyServerTask
__all__ = [ __all__ = [
"WebSocketProxyClientAgent", "WebSocketProxyClientTask",
"WebSocketProxyServerAgent", "WebSocketProxyServerTask",
] ]

131
src/pipecat/tasks/proxy/websocket/client.py
View File

@@ -10,7 +10,12 @@ import asyncio
from loguru import logger from loguru import logger
from pipecat.bus import BusMessage, BusTaskRegistryMessage, TaskBus from pipecat.bus import (
BusCancelTaskMessage,
BusEndTaskMessage,
BusMessage,
BusTaskRegistryMessage,
)
from pipecat.bus.messages import BusLocalMessage from pipecat.bus.messages import BusLocalMessage
from pipecat.bus.serializers import JSONMessageSerializer from pipecat.bus.serializers import JSONMessageSerializer
from pipecat.bus.serializers.base import MessageSerializer from pipecat.bus.serializers.base import MessageSerializer
@@ -22,17 +27,17 @@ try:
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
logger.error(f"Exception: {e}") logger.error(f"Exception: {e}")
logger.error( logger.error(
"In order to use WebSocketProxyClientAgent, you need to `pip install pipecat-ai-subagents[websocket]`." "In order to use WebSocketProxyClientTask, you need to `pip install pipecat-ai[websockets-base]`."
) )
raise Exception(f"Missing module: {e}") raise Exception(f"Missing module: {e}")
class WebSocketProxyClientAgent(BaseTask): class WebSocketProxyClientTask(BaseTask):
"""Forwards bus messages to a remote agent over WebSocket. """Forwards bus messages to a remote task over WebSocket.
Connects to a WebSocket URL and forwards messages between a local Connects to a WebSocket URL and forwards messages between a local
agent and a remote agent. Only messages targeted at the remote agent task and a remote task. Only messages targeted at the remote task
are sent. Only messages targeted at the local agent are accepted. are sent. Only messages targeted at the local task are accepted.
Event handlers available: Event handlers available:
@@ -41,61 +46,58 @@ class WebSocketProxyClientAgent(BaseTask):
Example:: Example::
proxy = WebSocketProxyClientAgent( proxy = WebSocketProxyClientTask(
"proxy", "proxy",
bus=runner.bus,
url="ws://remote-server:8765/ws", url="ws://remote-server:8765/ws",
remote_agent_name="worker", remote_task_name="worker",
local_agent_name="voice", local_task_name="voice",
) )
@proxy.event_handler("on_connected") @proxy.event_handler("on_connected")
async def on_connected(agent, websocket): async def on_connected(task, websocket):
logger.info("Connected to remote server") logger.info("Connected to remote server")
@proxy.event_handler("on_disconnected") @proxy.event_handler("on_disconnected")
async def on_disconnected(agent, websocket): async def on_disconnected(task, websocket):
logger.info("Disconnected from remote server") logger.info("Disconnected from remote server")
await runner.add_task(proxy) await runner.spawn(proxy)
""" """
def __init__( def __init__(
self, self,
name: str, name: str,
*, *,
bus: TaskBus,
url: str, url: str,
remote_agent_name: str, remote_task_name: str,
local_agent_name: str, local_task_name: str,
forward_messages: tuple[type[BusMessage], ...] = (), forward_messages: tuple[type[BusMessage], ...] = (),
headers: dict[str, str] | None = None, headers: dict[str, str] | None = None,
serializer: MessageSerializer | None = None, serializer: MessageSerializer | None = None,
): ):
"""Initialize the WebSocketProxyClientAgent. """Initialize the WebSocketProxyClientTask.
Args: Args:
name: Unique name for this agent. name: Unique name for this task.
bus: The `TaskBus` for inter-agent communication.
url: The WebSocket URL to connect to. url: The WebSocket URL to connect to.
remote_agent_name: Name of the agent on the remote server. remote_task_name: Name of the task on the remote server.
Only messages targeted at this agent are forwarded. Only messages targeted at this task are forwarded.
local_agent_name: Name of the local agent that should local_task_name: Name of the local task that should
receive responses. Only inbound messages targeted at receive responses. Only inbound messages targeted at
this agent are accepted. this task are accepted.
forward_messages: Additional message types to forward from forward_messages: Additional message types to forward from
the local agent (e.g. ``(BusFrameMessage,)`` for frame the local task (e.g. ``(BusFrameMessage,)`` for frame
routing). These are forwarded based on source agent name routing). These are forwarded based on source task name
only, regardless of target. only, regardless of target.
headers: Optional HTTP headers sent with the WebSocket headers: Optional HTTP headers sent with the WebSocket
handshake (e.g. for authentication). handshake (e.g. for authentication).
serializer: Serializer for bus messages. Defaults to serializer: Serializer for bus messages. Defaults to
`JSONMessageSerializer`. `JSONMessageSerializer`.
""" """
super().__init__(name, bus=bus) super().__init__(name)
self._url = url self._url = url
self._remote_agent_name = remote_agent_name self._remote_task_name = remote_task_name
self._local_agent_name = local_agent_name self._local_task_name = local_task_name
self._forward_messages = forward_messages self._forward_messages = forward_messages
self._headers = headers or {} self._headers = headers or {}
self._serializer = serializer or JSONMessageSerializer() self._serializer = serializer or JSONMessageSerializer()
@@ -105,22 +107,15 @@ class WebSocketProxyClientAgent(BaseTask):
self._register_event_handler("on_connected") self._register_event_handler("on_connected")
self._register_event_handler("on_disconnected") self._register_event_handler("on_disconnected")
async def cleanup(self):
"""Cancel the receive loop task and release resources."""
await super().cleanup()
if self._receive_task:
await self.cancel_task(self._receive_task)
self._receive_task = None
async def on_activated(self, args: dict | None) -> None: async def on_activated(self, args: dict | None) -> None:
"""Connect to the remote WebSocket server.""" """Connect to the remote WebSocket server."""
await super().on_activated(args) await super().on_activated(args)
logger.debug(f"Agent '{self}': connecting to {self._url}") logger.debug(f"Task '{self}': connecting to {self._url}")
self._ws = await connect(self._url, additional_headers=self._headers) self._ws = await connect(self._url, additional_headers=self._headers)
logger.debug(f"Agent '{self}': connected to {self._url}") logger.debug(f"Task '{self}': connected to {self._url}")
await self._call_event_handler("on_connected", self._ws) await self._call_event_handler("on_connected", self._ws)
@@ -129,8 +124,19 @@ class WebSocketProxyClientAgent(BaseTask):
# Schedule task right away. # Schedule task right away.
await asyncio.sleep(0) await asyncio.sleep(0)
async def stop(self) -> None:
"""Cancel the receive loop and close the WebSocket connection."""
if self._receive_task:
await self.cancel_task(self._receive_task)
self._receive_task = None
if self._ws:
await self._ws.close()
logger.debug(f"Task '{self}': WebSocket connection closed")
self._ws = None
await super().stop()
async def on_bus_message(self, message: BusMessage) -> None: async def on_bus_message(self, message: BusMessage) -> None:
"""Forward messages targeted at the remote agent. """Forward messages targeted at the remote task.
Args: Args:
message: The bus message to process. message: The bus message to process.
@@ -143,20 +149,23 @@ class WebSocketProxyClientAgent(BaseTask):
if isinstance(message, BusLocalMessage): if isinstance(message, BusLocalMessage):
return return
# Forward targeted messages to the remote agent # Forward targeted messages to the remote task.
if message.target == self._remote_agent_name: if message.target == self._remote_task_name:
await self._send_ws(message) await self._send_ws(message)
# Forward additional message types from the local agent # Forward additional message types from the local task.
elif isinstance(message, self._forward_messages): elif isinstance(message, self._forward_messages):
if message.source == self._local_agent_name: if message.source == self._local_task_name:
await self._send_ws(message) await self._send_ws(message)
async def _stop(self) -> None: async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
"""Close the WebSocket connection and stop.""" """Signal the run loop to finish on a graceful end."""
if self._ws: await super()._handle_task_end(message)
await self._ws.close() self._finished_event.set()
logger.debug(f"Agent '{self}': WebSocket connection closed")
await super()._stop() async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
"""Signal the run loop to finish on cancellation."""
await super()._handle_task_cancel(message)
self._finished_event.set()
async def _send_ws(self, message: BusMessage) -> None: async def _send_ws(self, message: BusMessage) -> None:
"""Serialize and send a message over the WebSocket.""" """Serialize and send a message over the WebSocket."""
@@ -165,9 +174,9 @@ class WebSocketProxyClientAgent(BaseTask):
try: try:
data = self._serializer.serialize(message) data = self._serializer.serialize(message)
await self._ws.send(data) await self._ws.send(data)
logger.trace(f"Agent '{self}': sent {message}") logger.trace(f"Task '{self}': sent {message}")
except websockets.exceptions.ConnectionClosed: except websockets.exceptions.ConnectionClosed:
logger.warning(f"Agent '{self}': connection closed, stopping forwarding") logger.warning(f"Task '{self}': connection closed, stopping forwarding")
ws = self._ws ws = self._ws
self._ws = None self._ws = None
await self._call_event_handler("on_disconnected", ws) await self._call_event_handler("on_disconnected", ws)
@@ -181,34 +190,36 @@ class WebSocketProxyClientAgent(BaseTask):
if not message: if not message:
continue continue
# Accept registry messages (target=None) for agent discovery # Accept registry messages (target=None) for task discovery.
if isinstance(message, BusTaskRegistryMessage): if isinstance(message, BusTaskRegistryMessage):
logger.trace( logger.trace(
f"Agent '{self}': received registry from remote: {message.agents}" f"Task '{self}': received registry from remote: {message.tasks}"
) )
await self.send_message(message) await self.send_message(message)
continue continue
# Accept additional message types (e.g. BusFrameMessage) # Accept additional message types (e.g. BusFrameMessage).
if self._forward_messages and isinstance(message, self._forward_messages): if self._forward_messages and isinstance(message, self._forward_messages):
logger.trace(f"Agent '{self}': received {message} from remote") logger.trace(f"Task '{self}': received {message} from remote")
await self.send_message(message) await self.send_message(message)
continue continue
# Only accept other messages targeted at the local agent # Only accept other messages targeted at the local task.
if message.target != self._local_agent_name: if message.target != self._local_task_name:
logger.warning( logger.warning(
f"Agent '{self}': dropped inbound message with " f"Task '{self}': dropped inbound message with "
f"unexpected target '{message.target}'" f"unexpected target '{message.target}'"
) )
continue continue
logger.trace(f"Agent '{self}': received {message} from remote") logger.trace(f"Task '{self}': received {message} from remote")
await self.send_message(message) await self.send_message(message)
except Exception: except Exception:
logger.exception(f"Agent '{self}': failed to deserialize remote message") logger.exception(f"Task '{self}': failed to deserialize remote message")
except websockets.exceptions.ConnectionClosed: except websockets.exceptions.ConnectionClosed:
logger.warning(f"Agent '{self}': WebSocket connection closed") logger.warning(f"Task '{self}': WebSocket connection closed")
ws = self._ws ws = self._ws
self._ws = None self._ws = None
await self._call_event_handler("on_disconnected", ws) await self._call_event_handler("on_disconnected", ws)
except asyncio.CancelledError:
pass

152
src/pipecat/tasks/proxy/websocket/server.py
View File

@@ -10,7 +10,12 @@ import asyncio
from loguru import logger from loguru import logger
from pipecat.bus import BusMessage, BusTaskRegistryMessage, TaskBus from pipecat.bus import (
BusCancelTaskMessage,
BusEndTaskMessage,
BusMessage,
BusTaskRegistryMessage,
)
from pipecat.bus.messages import BusLocalMessage from pipecat.bus.messages import BusLocalMessage
from pipecat.bus.serializers import JSONMessageSerializer from pipecat.bus.serializers import JSONMessageSerializer
from pipecat.bus.serializers.base import MessageSerializer from pipecat.bus.serializers.base import MessageSerializer
@@ -21,19 +26,17 @@ try:
from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState
except ModuleNotFoundError as e: except ModuleNotFoundError as e:
logger.error(f"Exception: {e}") logger.error(f"Exception: {e}")
logger.error( logger.error("In order to use WebSocketProxyServerTask, you need to `pip install starlette`.")
"In order to use WebSocketProxyServerAgent, you need to `pip install pipecat-ai-subagents[websocket]`."
)
raise Exception(f"Missing module: {e}") raise Exception(f"Missing module: {e}")
class WebSocketProxyServerAgent(BaseTask): class WebSocketProxyServerTask(BaseTask):
"""Receives bus messages from a remote client over WebSocket. """Receives bus messages from a remote client over WebSocket.
Accepts a FastAPI/Starlette WebSocket connection and forwards Accepts a FastAPI/Starlette WebSocket connection and forwards
messages between the remote client and a local agent. Only messages messages between the remote client and a local task. Only messages
from the local agent targeted at the remote agent are sent. Only from the local task targeted at the remote task are sent. Only
inbound messages targeted at the local agent are accepted. inbound messages targeted at the local task are accepted.
Event handlers available: Event handlers available:
@@ -45,58 +48,55 @@ class WebSocketProxyServerAgent(BaseTask):
@app.websocket("/ws") @app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket): async def websocket_endpoint(websocket: WebSocket):
await websocket.accept() await websocket.accept()
proxy = WebSocketProxyServerAgent( proxy = WebSocketProxyServerTask(
"gateway", "gateway",
bus=runner.bus,
websocket=websocket, websocket=websocket,
agent_name="worker", task_name="worker",
remote_agent_name="voice", remote_task_name="voice",
) )
@proxy.event_handler("on_client_connected") @proxy.event_handler("on_client_connected")
async def on_client_connected(agent, websocket): async def on_client_connected(task, websocket):
logger.info("Client connected") logger.info("Client connected")
@proxy.event_handler("on_client_disconnected") @proxy.event_handler("on_client_disconnected")
async def on_client_disconnected(agent, websocket): async def on_client_disconnected(task, websocket):
logger.info("Client disconnected") logger.info("Client disconnected")
await runner.add_task(proxy) await runner.spawn(proxy)
""" """
def __init__( def __init__(
self, self,
name: str, name: str,
*, *,
bus: TaskBus,
websocket: WebSocket, websocket: WebSocket,
agent_name: str, task_name: str,
remote_agent_name: str, remote_task_name: str,
forward_messages: tuple[type[BusMessage], ...] = (), forward_messages: tuple[type[BusMessage], ...] = (),
serializer: MessageSerializer | None = None, serializer: MessageSerializer | None = None,
): ):
"""Initialize the WebSocketProxyServerAgent. """Initialize the WebSocketProxyServerTask.
Args: Args:
name: Unique name for this agent. name: Unique name for this task.
bus: The `TaskBus` for inter-agent communication.
websocket: An accepted FastAPI/Starlette WebSocket connection. websocket: An accepted FastAPI/Starlette WebSocket connection.
agent_name: Name of the local agent to route messages to/from. task_name: Name of the local task to route messages to/from.
Only messages from this agent are forwarded to the client. Only messages from this task are forwarded to the client.
remote_agent_name: Name of the agent on the remote client. remote_task_name: Name of the task on the remote client.
Only outbound messages targeted at this agent are sent. Only outbound messages targeted at this task are sent.
Only inbound messages targeted at the local agent are accepted. Only inbound messages targeted at the local task are accepted.
forward_messages: Additional message types to forward from forward_messages: Additional message types to forward from
the local agent (e.g. ``(BusFrameMessage,)`` for frame the local task (e.g. ``(BusFrameMessage,)`` for frame
routing). These are forwarded based on source agent name routing). These are forwarded based on source task name
only, regardless of target. only, regardless of target.
serializer: Serializer for bus messages. Defaults to serializer: Serializer for bus messages. Defaults to
`JSONMessageSerializer`. `JSONMessageSerializer`.
""" """
super().__init__(name, bus=bus) super().__init__(name)
self._ws = websocket self._ws: WebSocket | None = websocket
self._agent_name = agent_name self._task_name = task_name
self._remote_agent_name = remote_agent_name self._remote_task_name = remote_task_name
self._forward_messages = forward_messages self._forward_messages = forward_messages
self._serializer = serializer or JSONMessageSerializer() self._serializer = serializer or JSONMessageSerializer()
self._receive_task: asyncio.Task | None = None self._receive_task: asyncio.Task | None = None
@@ -104,19 +104,11 @@ class WebSocketProxyServerAgent(BaseTask):
self._register_event_handler("on_client_connected") self._register_event_handler("on_client_connected")
self._register_event_handler("on_client_disconnected") self._register_event_handler("on_client_disconnected")
async def cleanup(self): async def start(self) -> None:
"""Cancel the receive loop task and release resources.""" """Start the WebSocket receive loop and watch the local task."""
await super().cleanup() await super().start()
if self._receive_task: logger.debug(f"Task '{self}': WebSocket proxy server ready")
await self.cancel_task(self._receive_task)
self._receive_task = None
async def on_ready(self) -> None:
"""Start receiving messages from the WebSocket and watch the local agent."""
await super().on_ready()
logger.debug(f"Agent '{self}': WebSocket proxy server ready")
await self._call_event_handler("on_client_connected", self._ws) await self._call_event_handler("on_client_connected", self._ws)
@@ -125,32 +117,41 @@ class WebSocketProxyServerAgent(BaseTask):
# Schedule task right away. # Schedule task right away.
await asyncio.sleep(0) await asyncio.sleep(0)
# Watch the local agent so we can notify the remote side when it's ready # Watch the local task so we can notify the remote side when it's ready.
await self.watch_task(self._agent_name) await self.watch_task(self._task_name)
async def stop(self) -> None:
"""Cancel the receive loop and close the WebSocket connection."""
if self._receive_task:
await self.cancel_task(self._receive_task)
self._receive_task = None
if self._ws and self._ws.client_state == WebSocketState.CONNECTED:
await self._ws.close()
logger.debug(f"Task '{self}': WebSocket connection closed")
await super().stop()
async def on_task_ready(self, data: TaskReadyData) -> None: async def on_task_ready(self, data: TaskReadyData) -> None:
"""Notify the remote client that the local agent is ready.""" """Notify the remote client that the local task is ready."""
if not self._ws: if not self._ws:
return return
if data.agent_name != self._agent_name: if data.task_name != self._task_name:
return return
logger.debug(f"Agent '{self}': local agent '{self._agent_name}' ready, notifying remote") logger.debug(f"Task '{self}': local task '{self._task_name}' ready, notifying remote")
try: try:
msg = BusTaskRegistryMessage( msg = BusTaskRegistryMessage(
source=self.name, source=self.name,
runner=data.runner, runner=data.runner,
agents=[TaskRegistryEntry(name=self._agent_name)], tasks=[TaskRegistryEntry(name=self._task_name)],
) )
await self._send_ws(msg) await self._send_ws(msg)
except Exception: except Exception:
logger.exception(f"Agent '{self}': failed to send registry to remote") logger.exception(f"Task '{self}': failed to send registry to remote")
async def on_bus_message(self, message: BusMessage) -> None: async def on_bus_message(self, message: BusMessage) -> None:
"""Forward messages from the local agent to the remote client. """Forward messages from the local task to the remote client.
Args: Args:
message: The bus message to process. message: The bus message to process.
@@ -163,22 +164,25 @@ class WebSocketProxyServerAgent(BaseTask):
if isinstance(message, BusLocalMessage): if isinstance(message, BusLocalMessage):
return return
if message.source != self._agent_name: if message.source != self._task_name:
return return
# Forward targeted messages from the local agent to the remote agent # Forward targeted messages from the local task to the remote task.
if message.target == self._remote_agent_name: if message.target == self._remote_task_name:
await self._send_ws(message) await self._send_ws(message)
# Forward additional message types from the local agent # Forward additional message types from the local task.
elif isinstance(message, self._forward_messages): elif isinstance(message, self._forward_messages):
await self._send_ws(message) await self._send_ws(message)
async def _stop(self) -> None: async def _handle_task_end(self, message: BusEndTaskMessage) -> None:
"""Close the WebSocket connection and stop.""" """Signal the run loop to finish on a graceful end."""
if self._ws and self._ws.client_state == WebSocketState.CONNECTED: await super()._handle_task_end(message)
await self._ws.close() self._finished_event.set()
logger.debug(f"Agent '{self}': WebSocket connection closed")
await super()._stop() async def _handle_task_cancel(self, message: BusCancelTaskMessage) -> None:
"""Signal the run loop to finish on cancellation."""
await super()._handle_task_cancel(message)
self._finished_event.set()
async def _send_ws(self, message: BusMessage) -> None: async def _send_ws(self, message: BusMessage) -> None:
"""Serialize and send a message over the WebSocket.""" """Serialize and send a message over the WebSocket."""
@@ -187,9 +191,9 @@ class WebSocketProxyServerAgent(BaseTask):
try: try:
data = self._serializer.serialize(message) data = self._serializer.serialize(message)
await self._ws.send_bytes(data) await self._ws.send_bytes(data)
logger.trace(f"Agent '{self}': sent {message}") logger.trace(f"Task '{self}': sent {message}")
except (WebSocketDisconnect, Exception): except WebSocketDisconnect:
logger.warning(f"Agent '{self}': connection closed, stopping forwarding") logger.warning(f"Task '{self}': connection closed, stopping forwarding")
ws = self._ws ws = self._ws
self._ws = None self._ws = None
await self._call_event_handler("on_client_disconnected", ws) await self._call_event_handler("on_client_disconnected", ws)
@@ -204,26 +208,26 @@ class WebSocketProxyServerAgent(BaseTask):
if not message: if not message:
continue continue
# Accept additional message types (e.g. BusFrameMessage) # Accept additional message types (e.g. BusFrameMessage).
if self._forward_messages and isinstance(message, self._forward_messages): if self._forward_messages and isinstance(message, self._forward_messages):
logger.trace(f"Agent '{self}': received {message} from client") logger.trace(f"Task '{self}': received {message} from client")
await self.send_message(message) await self.send_message(message)
continue continue
# Only accept other messages targeted at the local agent # Only accept other messages targeted at the local task.
if message.target != self._agent_name: if message.target != self._task_name:
logger.warning( logger.warning(
f"Agent '{self}': dropped inbound message with " f"Task '{self}': dropped inbound message with "
f"unexpected target '{message.target}'" f"unexpected target '{message.target}'"
) )
continue continue
logger.trace(f"Agent '{self}': received {message} from client") logger.trace(f"Task '{self}': received {message} from client")
await self.send_message(message) await self.send_message(message)
except Exception: except Exception:
logger.exception(f"Agent '{self}': failed to deserialize client message") logger.exception(f"Task '{self}': failed to deserialize client message")
except WebSocketDisconnect: except WebSocketDisconnect:
logger.warning(f"Agent '{self}': client disconnected") logger.warning(f"Task '{self}': client disconnected")
ws = self._ws ws = self._ws
self._ws = None self._ws = None
await self._call_event_handler("on_client_disconnected", ws) await self._call_event_handler("on_client_disconnected", ws)