Basic web loop works

This commit is contained in:
Niels Geens
2024-11-07 15:10:17 +01:00
parent 6e78a0f0bd
commit 970414e56d
9 changed files with 218 additions and 132 deletions

View File

@@ -34,8 +34,10 @@ class TestLLM:
yield "</reasoning>" yield "</reasoning>"
class Main: class Main:
def __init__(self): @classmethod
self._config = Config() async def create(cls, config: Config):
self = cls()
self._config = config
self._system_prompt = self._config.system_prompt.read_text() self._system_prompt = self._config.system_prompt.read_text()
self._action_schema = self._config.action_schema.read_text() self._action_schema = self._config.action_schema.read_text()
@@ -62,13 +64,14 @@ class Main:
validator=XMLValidator(self._action_schema), validator=XMLValidator(self._action_schema),
parser=ResponseParser(self._io_buffer) parser=ResponseParser(self._io_buffer)
) )
self._ws_manager = WebSocketManager( self._ws_manager = await WebSocketManager.create(
self._agent, self._agent,
self._io_buffer self._io_buffer
) )
self._app = web.Application() self._app = web.Application()
self._init_routes() self._init_routes()
return self
@property @property
def app(self): def app(self):
@@ -114,6 +117,8 @@ class Main:
return response return response
if __name__ == "__main__": if __name__ == "__main__":
main = Main() loop = asyncio.new_event_loop()
print("Web server started at http://localhost:8080") config = Config()
web.run_app(main.app, port=8080) main = loop.run_until_complete(Main.create(config))
print(f"Web server started at http://{config.host}:{config.port}")
web.run_app(main.app, loop=loop, host=config.host, port=config.port)

View File

@@ -42,8 +42,8 @@ class Config:
parser.add_argument( parser.add_argument(
'--host', '--host',
type=str, type=str,
default=os.getenv('SIA_SERVER_HOST', 'localhost'), default=os.getenv('SIA_SERVER_HOST', '0.0.0.0'),
help='Web server host (default: localhost, env: SIA_SERVER_HOST)' help='Web server host (default: 0.0.0.0, env: SIA_SERVER_HOST)'
) )
parser.add_argument( parser.add_argument(
'--port', '--port',

View File

@@ -1,4 +1,5 @@
from enum import Enum, auto from enum import Enum, auto
from threading import Thread, Lock
from typing import Callable, List, Optional from typing import Callable, List, Optional
from .base_agent import BaseAgent from .base_agent import BaseAgent
@@ -13,17 +14,15 @@ from .xml_validator import XMLValidator
class WebAgentState(Enum): class WebAgentState(Enum):
""" """
States for the web agent state machine. States for the web agent state machine.
States:
UPDATE: Updating system metrics and working memory entries
CONTEXT_APPROVAL: Waiting for human approval of context
INFERENCE: Processing context through LLM
RESPONSE_APPROVAL: Waiting for human approval of LLM response
""" """
UPDATE = auto() UPDATE = auto()
"""Updating system metrics and working memory entries"""
CONTEXT_APPROVAL = auto() CONTEXT_APPROVAL = auto()
"""Waiting for human approval of context"""
INFERENCE = auto() INFERENCE = auto()
"""Processing context through LLM"""
RESPONSE_APPROVAL = auto() RESPONSE_APPROVAL = auto()
"""Waiting for human approval of LLM response"""
STOPPED = auto() STOPPED = auto()
class WebAgent(BaseAgent): class WebAgent(BaseAgent):
@@ -48,17 +47,19 @@ class WebAgent(BaseAgent):
super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser) super().__init__(action_schema, working_memory, system_metrics, llm, validator, parser)
self._response = "" self._response = ""
self._system_prompt = system_prompt self._system_prompt = system_prompt
self._current_state = WebAgentState.CONTEXT_APPROVAL self._state = WebAgentState.CONTEXT_APPROVAL
self._validation_error: Optional[str] = None self._validation_error: Optional[str] = None
self._state_change_handlers: List[Callable[[WebAgentState], None]] = [] self._state_change_handlers: List[Callable[[WebAgentState], None]] = []
self._context_change_handlers: List[Callable[[str], None]] = []
self._response_change_handlers: List[Callable[[str, str], None]] = [] self._response_change_handlers: List[Callable[[str, str], None]] = []
self._command_result: Optional[CommandResult] = None self._command_result: Optional[CommandResult] = None
self.context = self._compile_context() self._context = self._compile_context()
self._state_lock = Lock()
@property @property
def current_state(self) -> WebAgentState: def state(self) -> WebAgentState:
"""Get the current state of the agent.""" """Get the current state of the agent."""
return self._current_state return self._state
@property @property
def command_result(self) -> Optional[CommandResult]: def command_result(self) -> Optional[CommandResult]:
@@ -75,6 +76,11 @@ class WebAgent(BaseAgent):
"""Get the current response.""" """Get the current response."""
return self._response return self._response
@property
def context(self) -> str:
"""Get the current context."""
return self._context
def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None: def add_state_change_handler(self, handler: Callable[[WebAgentState], None]) -> None:
""" """
Add a callback for state changes. Add a callback for state changes.
@@ -85,6 +91,16 @@ class WebAgent(BaseAgent):
if handler not in self._state_change_handlers: if handler not in self._state_change_handlers:
self._state_change_handlers.append(handler) self._state_change_handlers.append(handler)
def add_context_change_handler(self, handler: Callable[[str], None]) -> None:
"""
Add a callback for context changes.
Args:
handler: Function to call with new context
"""
if handler not in self._context_change_handlers:
self._context_change_handlers.append(handler)
def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None: def add_response_change_handler(self, handler: Callable[[str, str], None]) -> None:
""" """
Add a callback for response changes. Add a callback for response changes.
@@ -95,52 +111,53 @@ class WebAgent(BaseAgent):
if handler not in self._response_change_handlers: if handler not in self._response_change_handlers:
self._response_change_handlers.append(handler) self._response_change_handlers.append(handler)
def _notify_state_change(self) -> None: def modify_context(self, context: str) -> None:
"""Notify all handlers of state change.""" with self._state_lock:
for handler in self._state_change_handlers: if self._state != WebAgentState.CONTEXT_APPROVAL:
handler(self._current_state) error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
raise Exception(error_msg)
Thread(target=self._set_context, args=(context,)).start()
def set_response(self, response: str) -> None: def modify_response(self, response: str) -> None:
"""Set response and notify handlers.""" started = False
self._response = response try:
validation_error = self._validator.validate(response) self._state_lock.acquire()
self._validation_error = validation_error if self._state != WebAgentState.RESPONSE_APPROVAL:
for handler in self._response_change_handlers: error_msg = f"Not in RESPONSE_APPROVAL state (current: {self._state})"
handler(response, validation_error) raise Exception(error_msg)
Thread(target=self._set_response, args=(response, self._state_lock)).start()
def _set_state(self, new_state: WebAgentState) -> None: started = True
""" finally:
Set the current state and notify handlers. if not started:
self._state_lock.release()
Args:
new_state: New state to set
"""
self._current_state = new_state
self._notify_state_change()
def approve_context(self) -> None: def approve_context(self) -> None:
if self._current_state != WebAgentState.CONTEXT_APPROVAL: with self._state_lock:
error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._current_state})" if self._state != WebAgentState.CONTEXT_APPROVAL:
raise Exception(error_msg) error_msg = f"Not in CONTEXT_APPROVAL state (current: {self._state})"
raise Exception(error_msg)
self._set_state(WebAgentState.INFERENCE)
Thread(target=self._approve_context_thread).start()
self._set_state(WebAgentState.INFERENCE) def approve_response(self) -> None:
self.set_response("") with self._state_lock:
if self._state != WebAgentState.RESPONSE_APPROVAL:
raise Exception("Not in RESPONSE_APPROVAL state")
self._set_state(WebAgentState.UPDATE)
Thread(target=self._approve_response_thread).start()
def _approve_context_thread(self) -> None:
self._set_response("")
response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context) response_token_iter = self._llm.infer(f"{self._system_prompt}\n{self._action_schema}", self.context)
response = "" response = ""
for token in response_token_iter: for token in response_token_iter:
response += token response += token
self.set_response(response) self._set_response(response)
print(f"{token}") print(f"{token}")
self._set_state(WebAgentState.RESPONSE_APPROVAL) self._set_state(WebAgentState.RESPONSE_APPROVAL)
def approve_response(self) -> None: def _approve_response_thread(self) -> None:
if self._current_state != WebAgentState.RESPONSE_APPROVAL: parse_result = self._parser.parse(self._response)
raise Exception("Not in RESPONSE_APPROVAL state")
self._set_state(WebAgentState.UPDATE)
parse_result = self._parser.parse(self.response)
if isinstance(parse_result, Command): if isinstance(parse_result, Command):
result = parse_result.execute(self._working_memory) result = parse_result.execute(self._working_memory)
self._command_result = result self._command_result = result
@@ -151,5 +168,29 @@ class WebAgent(BaseAgent):
self._working_memory.add_entry(parse_result) self._working_memory.add_entry(parse_result)
self._working_memory.update() self._working_memory.update()
self.context = self._compile_context() self._set_context(self._compile_context())
self._set_state(WebAgentState.CONTEXT_APPROVAL) self._set_state(WebAgentState.CONTEXT_APPROVAL)
def _set_state(self, state) -> None:
"""Notify all handlers of state change."""
self._state = state
for handler in self._state_change_handlers:
handler(self._state)
def _set_context(self, context) -> None:
"""Notify all handlers of context change."""
self._context = context
for handler in self._context_change_handlers:
handler(context)
def _set_response(self, response: str, lock: Optional[Lock] = None) -> None:
"""Set response and notify handlers."""
try:
self._response = response
validation_error = self._validator.validate(response)
self._validation_error = validation_error
for handler in self._response_change_handlers:
handler(response, validation_error)
finally:
if lock is not None:
lock.release()

View File

@@ -27,18 +27,23 @@ class WebSocketManager:
""" """
Manages WebSocket connections and integrates with the WebAgent Manages WebSocket connections and integrates with the WebAgent
""" """
def __init__(self, @classmethod
agent: WebAgent, async def create(
io_buffer: WebIOBuffer): cls,
agent: WebAgent,
io_buffer: WebIOBuffer):
"""Initialize the web server.""" """Initialize the web server."""
self = cls()
self.agent = agent self.agent = agent
self.io_buffer = io_buffer self.io_buffer = io_buffer
self._clients: Set[web.WebSocketResponse] = set() self._clients: Set[web.WebSocketResponse] = set()
self.agent.add_state_change_handler(self._handle_state_change) self.agent.add_state_change_handler(self._wrap_async(self._handle_state_change))
self.agent.add_response_change_handler(self._handle_response_change) self.agent.add_context_change_handler(self._wrap_async(self._handle_context_change))
self.agent.add_response_change_handler(self._wrap_async(self._handle_response_change))
return self
async def broadcast_message(self, message: Dict): async def _broadcast_message(self, message: Dict):
"""Broadcast message to all connected clients.""" """Broadcast message to all connected clients."""
disconnected = set() disconnected = set()
for ws in self._clients: for ws in self._clients:
@@ -48,20 +53,62 @@ class WebSocketManager:
disconnected.add(ws) disconnected.add(ws)
self._clients -= disconnected self._clients -= disconnected
def _handle_state_change(self, new_state: WebAgentState): def _wrap_async(self, coro_func):
"""
Wraps an async callback to be safely called from another thread.
Args:
coro_func: The async function to wrap
Returns:
A sync function that schedules the coroutine on the event loop
"""
loop = asyncio.get_event_loop()
def wrapper(*args, **kwargs):
loop.call_soon_threadsafe(lambda: asyncio.create_task(coro_func(*args, **kwargs)))
return wrapper
async def _handle_state_change(self, new_state: WebAgentState):
"""Handle state changes from the WebAgent.""" """Handle state changes from the WebAgent."""
asyncio.run(self.broadcast_message({ await self._broadcast_message({
"type": ServerMessage.STATE_CHANGE, "type": ServerMessage.STATE_CHANGE,
"state": new_state.name, "state": new_state.name,
})) })
def _handle_response_change(self, response: str, validation_error: str): async def _handle_context_change(self, context: str):
"""Handle context changes from the WebAgent."""
await self._broadcast_message({
"type": ServerMessage.CONTEXT_UPDATE,
"context": context
})
async def _handle_response_change(self, response: str, validation_error: str):
"""Handle response changes from the WebAgent.""" """Handle response changes from the WebAgent."""
asyncio.run(self.broadcast_message({ await self._broadcast_message({
"type": ServerMessage.RESPONSE_UPDATE, "type": ServerMessage.RESPONSE_UPDATE,
"response": response, "response": response,
"validation_error": validation_error "validation_error": validation_error
})) })
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
"""Handle incoming client message."""
message_type = data.get("type")
if message_type == ClientMessage.APPROVE_CONTEXT:
if self.agent.state == WebAgentState.CONTEXT_APPROVAL:
if data.get("context") is not None:
self.agent.modify_context(data.get("context"))
self.agent.approve_context()
elif message_type == ClientMessage.APPROVE_RESPONSE:
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
if data.get("response") is not None:
self.agent.modify_response(data.get("response"))
self.agent.approve_response()
elif message_type == ClientMessage.MODIFY_RESPONSE:
if self.agent.state == WebAgentState.RESPONSE_APPROVAL:
self.agent.modify_response(data.get("response"))
elif message_type == ClientMessage.SEND_INPUT:
self.io_buffer.append_stdin(data.get("input"))
async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse: async def handle_websocket(self, request: web.Request) -> web.WebSocketResponse:
"""Handle new WebSocket connections and messages.""" """Handle new WebSocket connections and messages."""
@@ -73,7 +120,7 @@ class WebSocketManager:
try: try:
await ws.send_json({ await ws.send_json({
"type": ServerMessage.STATE_CHANGE, "type": ServerMessage.STATE_CHANGE,
"state": self.agent.current_state.name "state": self.agent.state.name
}) })
await ws.send_json({ await ws.send_json({
"type": ServerMessage.CONTEXT_UPDATE, "type": ServerMessage.CONTEXT_UPDATE,
@@ -93,19 +140,3 @@ class WebSocketManager:
finally: finally:
self._clients.remove(ws) self._clients.remove(ws)
return ws return ws
async def _handle_client_message(self, request: web.Request, data: dict[str, Any]):
"""Handle incoming client message."""
message_type = data.get("type")
if message_type == ClientMessage.APPROVE_CONTEXT:
if self.agent.current_state == WebAgentState.CONTEXT_APPROVAL:
await asyncio.to_thread(self.agent.approve_context)
elif message_type == ClientMessage.APPROVE_RESPONSE:
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
await asyncio.to_thread(self.agent.approve_response)
elif message_type == ClientMessage.MODIFY_RESPONSE:
if self.agent.current_state == WebAgentState.RESPONSE_APPROVAL:
await asyncio.to_thread(self.agent.set_response, data.get("response"))
elif message_type == ClientMessage.SEND_INPUT:
self.io_buffer.append_stdin(data.get("input"))

View File

@@ -94,7 +94,7 @@ class WebAgentTest(unittest.TestCase):
def test_initialization(self): def test_initialization(self):
"""Test initial state of web agent.""" """Test initial state of web agent."""
self.assertEqual(self.agent.current_state, WebAgentState.CONTEXT_APPROVAL) self.assertEqual(self.agent.state, WebAgentState.CONTEXT_APPROVAL)
self.assertEqual(self.agent.response, "") self.assertEqual(self.agent.response, "")
self.assertIsNotNone(self.agent.context) self.assertIsNotNone(self.agent.context)
self.assertIsNone(self.agent.validation_error) self.assertIsNone(self.agent.validation_error)
@@ -121,7 +121,7 @@ class WebAgentTest(unittest.TestCase):
def test_approve_context_state_error(self): def test_approve_context_state_error(self):
"""Test approving context in wrong state.""" """Test approving context in wrong state."""
# Force agent into UPDATE state # Force agent into UPDATE state
self.agent._current_state = WebAgentState.UPDATE self.agent._state = WebAgentState.UPDATE
with self.assertRaises(Exception) as context: with self.assertRaises(Exception) as context:
asyncio.run(self.agent.approve_context()) asyncio.run(self.agent.approve_context())
@@ -169,7 +169,7 @@ class WebAgentTest(unittest.TestCase):
self.parser.parse = Mock(return_value=mock_command) self.parser.parse = Mock(return_value=mock_command)
# Set initial state and response # Set initial state and response
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL self.agent._state = WebAgentState.RESPONSE_APPROVAL
self.agent._response = "<stop/>" self.agent._response = "<stop/>"
# Approve response # Approve response
@@ -194,8 +194,8 @@ class WebAgentTest(unittest.TestCase):
self.agent.add_state_change_handler(self.state_change_handler) self.agent.add_state_change_handler(self.state_change_handler)
# Set initial state and response # Set initial state and response
self.agent._current_state = WebAgentState.RESPONSE_APPROVAL self.agent._state = WebAgentState.RESPONSE_APPROVAL
self.agent.set_response("<reasoning>test reasoning</reasoning>") self.agent._set_response("<reasoning>test reasoning</reasoning>")
# Approve response # Approve response
self.agent.approve_response() self.agent.approve_response()
@@ -218,7 +218,7 @@ class WebAgentTest(unittest.TestCase):
self.agent.add_response_change_handler(self.response_change_handler) self.agent.add_response_change_handler(self.response_change_handler)
error_message = "Invalid XML" error_message = "Invalid XML"
self.mock_validator.validate.return_value = error_message self.mock_validator.validate.return_value = error_message
self.agent.set_response("<invalid>") self.agent._set_response("<invalid>")
self.assertEqual(self.agent.validation_error, error_message) self.assertEqual(self.agent.validation_error, error_message)
self.assertEqual(self.response_changes, ["<invalid>"]) self.assertEqual(self.response_changes, ["<invalid>"])

View File

@@ -61,7 +61,7 @@ class WebSocketManagerTest(AioHTTPTestCase):
) )
# Create WebSocketManager # Create WebSocketManager
self.ws_manager = WebSocketManager(self.agent, self.io_buffer) self.ws_manager = await WebSocketManager.create(self.agent, self.io_buffer)
# Create application # Create application
app = web.Application() app = web.Application()

View File

@@ -24,35 +24,41 @@ const App = () => {
const [showSidebar, setShowSidebar] = useState(false); const [showSidebar, setShowSidebar] = useState(false);
const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL); const [agentState, setAgentState] = useState(AgentState.CONTEXT_APPROVAL);
const { wsState, sendMessage, setMessageHandler } = useWebSocket('ws://localhost:8080/ws'); const { wsState, sendMessage, addMessageHandler } = useWebSocket('ws://localhost:8080/ws');
// Set up WebSocket message handlers
useEffect(() => { useEffect(() => {
setMessageHandler(ServerMessageType.STATE_CHANGE, (message) => { addMessageHandler((message) => {
setAgentState(message.state); switch(message.type) {
case ServerMessageType.STATE_CHANGE:
setAgentState(message.state);
break;
case ServerMessageType.CONTEXT_UPDATE:
setOriginalContext(message.context);
setModifiedContext(message.context);
break;
case ServerMessageType.RESPONSE_UPDATE:
setOriginalResponse(message.response);
setModifiedResponse(message.response);
setValidationError(message.validation_error);
break;
case ServerMessageType.OUTPUT_UPDATE:
setOutput(prev => prev + message.data);
break;
}
}); });
}, [addMessageHandler]);
setMessageHandler(ServerMessageType.CONTEXT_UPDATE, (message) => {
setOriginalContext(message.context);
setModifiedContext(message.context);
});
setMessageHandler(ServerMessageType.RESPONSE_UPDATE, (message) => {
setOriginalResponse(message.response);
setModifiedResponse(message.response);
setValidationError(message.validation_error);
});
setMessageHandler(ServerMessageType.OUTPUT_UPDATE, (message) => {
setOutput(prev => prev + message.data);
});
}, [setMessageHandler]);
const handleApprove = () => { const handleApprove = () => {
if (agentState === AgentState.CONTEXT_APPROVAL) { if (agentState === AgentState.CONTEXT_APPROVAL) {
sendMessage(ClientMessageType.APPROVE_CONTEXT); sendMessage({
"type": ClientMessageType.APPROVE_CONTEXT,
"context": modifiedContext
});
} else if (agentState === AgentState.RESPONSE_APPROVAL) { } else if (agentState === AgentState.RESPONSE_APPROVAL) {
sendMessage(ClientMessageType.APPROVE_RESPONSE, modifiedResponse); sendMessage({
"type": ClientMessageType.APPROVE_RESPONSE,
"response": modifiedResponse
});
} }
}; };
@@ -63,15 +69,28 @@ const App = () => {
} }
}; };
useEffect(() => {
if (agentState === AgentState.APPROVE_RESPONSE) {
sendMessage({
"type": ClientMessageType.MODIFY_RESPONSE,
"response": modifiedResponse
});
}
}, [modifiedResponse]);
useEffect(() => { useEffect(() => {
switch (agentState) { switch (agentState) {
case AgentState.UPDATE: // fall through
case AgentState.CONTEXT_APPROVAL: case AgentState.CONTEXT_APPROVAL:
setActiveTab(Tabs.CONTEXT); setActiveTab(Tabs.CONTEXT);
break; break;
case AgentState.INFERENCE: // fall through
case AgentState.RESPONSE_APPROVAL: case AgentState.RESPONSE_APPROVAL:
setActiveTab(Tabs.RESPONSE); setActiveTab(Tabs.RESPONSE);
break; break;
} }
setShowDiff(false);
setShowSidebar(false);
}, [agentState]); }, [agentState]);
return ( return (

View File

@@ -19,7 +19,6 @@ export const Tabs = {
}; };
export const ClientMessageType = { export const ClientMessageType = {
MODIFY_CONTEXT: 'MODIFY_CONTEXT',
APPROVE_CONTEXT: 'APPROVE_CONTEXT', APPROVE_CONTEXT: 'APPROVE_CONTEXT',
MODIFY_RESPONSE: 'MODIFY_RESPONSE', MODIFY_RESPONSE: 'MODIFY_RESPONSE',
APPROVE_RESPONSE: 'APPROVE_RESPONSE', APPROVE_RESPONSE: 'APPROVE_RESPONSE',

View File

@@ -4,7 +4,7 @@ import { WebSocketState, MessageType } from '../constants';
export const useWebSocket = (url) => { export const useWebSocket = (url) => {
const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED); const [wsState, setWsState] = useState(WebSocketState.DISCONNECTED);
const wsRef = useRef(null); const wsRef = useRef(null);
const messageHandlersRef = useRef({}); const handlersRef = useRef([]);
const connect = useCallback(() => { const connect = useCallback(() => {
setWsState(WebSocketState.CONNECTING); setWsState(WebSocketState.CONNECTING);
@@ -27,12 +27,7 @@ export const useWebSocket = (url) => {
ws.onmessage = (event) => { ws.onmessage = (event) => {
try { try {
const message = JSON.parse(event.data); const message = JSON.parse(event.data);
const handler = messageHandlersRef.current[message.type]; handlersRef.current.forEach(handler => handler(message));
if (handler) {
handler(message);
} else {
console.warn('No handler for message type:', message.type);
}
} catch (error) { } catch (error) {
console.error('Error handling WebSocket message:', error); console.error('Error handling WebSocket message:', error);
} }
@@ -48,25 +43,21 @@ export const useWebSocket = (url) => {
}; };
}, [connect]); }, [connect]);
const sendMessage = useCallback((type, data = null) => { const sendMessage = useCallback((data) => {
if (wsRef.current?.readyState === WebSocket.OPEN) { if (wsRef.current?.readyState === WebSocket.OPEN) {
const message = { wsRef.current.send(JSON.stringify(data));
type,
data: data || undefined
};
wsRef.current.send(JSON.stringify(message));
} else { } else {
console.error('WebSocket not connected, state:', wsRef.current?.readyState); console.error('WebSocket not connected, state:', wsRef.current?.readyState);
} }
}, []); }, []);
const setMessageHandler = useCallback((type, handler) => { const addMessageHandler = useCallback((handler) => {
messageHandlersRef.current[type] = handler; handlersRef.current.push(handler);
}, []); }, []);
return { return {
wsState, wsState,
sendMessage, sendMessage,
setMessageHandler addMessageHandler
}; };
}; };