diff --git a/backend/src/server/websocket.py b/backend/src/server/websocket.py index 0887453..d7e45b1 100644 --- a/backend/src/server/websocket.py +++ b/backend/src/server/websocket.py @@ -68,275 +68,10 @@ def register(server, app, auth_manager): stop_listening = app.on_event(on_event) while True: - req = json.loads(await websocket.receive()) - if not 'action' in req: - await send({'error': 'no action'}) - continue - action = req['action'] - body = req.get('body', {}) - nonce = req.get('nonce', None) - print(f"[ws] action @{user_id} {action}") + asyncio.create_task(handle_request( + app, user_id, session_connections, send, await websocket.receive() + )) - async def agent_id_from_body (allow_none=False): - if 'agentId' in body: - return body['agentId'] - elif 'agentname' in body: - return await app.agent_id_from_name(body['agentname']) - elif allow_none: - return None - else: - raise Exception("field 'agentId' or 'agentname' required") - - if action == "ping": - await send({'message': "pong"}) - - elif action == "create-agent": - name = body['name'] - agent_id = await app.create_agent( - user_id=user_id, - name=name, - agentname=body.get('agentname'), - description=body.get('description'), - picture=body.get('picture'), - kind=body['kind'], - data=body.get('data'), - temporary=body.get('temporary'), - ) - await send({'message': "agent created", 'agentId': agent_id, 'nonce': nonce}) - - elif action == "update-agent": - agent_id = await agent_id_from_body() - await app.update_agent( - agent_id=agent_id, - name=body.get('name'), - agentname=body.get('agentname'), - description=body.get('description'), - picture=body.get('picture'), - data=body.get('data', {}) - ) - await send({'message': "agent updated", 'agentId': agent_id, 'nonce': nonce}) - - elif action == "get-models": - await send({ - 'models': [{"id": model.id, "name": model.name, "value": model.value, "kind": model.kind} for model in await app.get_models()], - 'nonce': nonce - }) - - elif action == "save-model": - model = await app.save_model( - model_id=body.get("modelId"), - name=body.get("name"), - kind=body.get("kind"), - value=body.get("value"), - ) - if model: - await send({'message': "model saved", 'modelId': model.id, 'nonce': nonce}) - elif "modelId" in body: - await send({'message': "model deleted", 'modelId': body["modelId"], 'nonce': nonce}) - - elif action == "get-agent-config-fields": - await send({'fields': await app.get_agent_config_fields(await agent_id_from_body()), 'nonce': nonce}) - - elif action == "get-agent-kinds": - await send({'message': "agents retrieved", 'kinds': app.get_agent_kinds(), 'nonce': nonce}) - - elif action == "get-agent-kind-info": - await send({'message': "agent kind retrieved", 'info': app.get_agent_kind_info(body['kind']), 'nonce': nonce}) - - elif action == "get-agents": - await send({'message': "agents retrieved", 'agents': await app.get_agents(user_id), 'nonce': nonce}) - - elif action == "get-agent-resources": - values = await app.get_agent_resources(await agent_id_from_body(), body.get('kind'), body.get('key')) - await send({'resources': [{'kind': r.kind, 'key': r.key, 'value': r.value} for r in values], 'nonce': nonce}) - - elif action == "save-agent-resource": - await app.save_agent_resource(await agent_id_from_body(), body['kind'], body['key'], body['value']) - await send({'message': "resource saved", 'nonce': nonce}) - - elif action == "delete-agent-resource": - await app.delete_agent_resource(await agent_id_from_body(), body['kind'], body['key']) - await send({'message': "resource deleted", 'nonce': nonce}) - - elif action == "get-agent-runs": - await send({ - 'runs': await app.get_agent_runs(agent_id=await agent_id_from_body(allow_none=True), user_id=user_id), - 'nonce': nonce - }) - - elif action == "cancel-agent-run": - await app.cancel_agent_run(body['runId']) - - elif action == "remote-list": - await send({'remotes': app.listRemotes(user_id), 'nonce': nonce}) - - elif action == "remote-connect": - # Connect a remote served by THIS session's channel. - # Expected failures (duplicate name) are replied, not - # raised, so the session stays alive. - try: - connection = app.connectRemote(user_id, body['remotename'], body.get('functions', []), send) - session_connections[body['remotename']] = connection - await send({'message': "remote connected", 'remotename': body['remotename'], 'nonce': nonce}) - except Exception as e: - await send({'error': str(e), 'nonce': nonce}) - - elif action == "remote-disconnect": - # Manual disconnect, allowed from any session of the user - connection = app.disconnectRemote(user_id, body['remotename']) - session_connections.pop(body['remotename'], None) - if connection is None: - await send({'error': "no connected remote named '" + str(body['remotename']) + "'", 'nonce': nonce}) - else: - await send({'message': "remote disconnected", 'remotename': body['remotename'], 'nonce': nonce}) - - elif action == "remote-call": - # Invoke a function of one of the user's connected remotes. - # Expected failures (unknown remote/function, timeout, - # remote-side error) are replied, not raised. - - - async def _handle_remote_call(user_id, remotename, func_name, args, send, nonce): - try: - result = await app.callRemote(user_id, remotename, func_name, args) - await send({'result': result, 'nonce': nonce}) - except Exception as e: - await send({'error': str(e), 'nonce': nonce}) - asyncio.create_task(_handle_remote_call( - user_id, body['remotename'], body['function'], body.get('args', {}), send, nonce - )) - - elif action == "remote-call-result": - # A session serving remotes replies to a downstream call. - # Fields live at the message top level, not in 'body'. - # requestIds are '#' and resolve/reject - # are no-ops for unknown ids, so fanning out over this - # session's connections routes the reply correctly. - app.resolve_remote_call(body['requestId'], body.get('result'), error=body.get('error')) - - elif action == "create-chat": - chat_id = await app.create_chat( - user_id=user_id, - name=body['name'], - description=body.get('description'), - picture=body.get('picture'), - members=body.get('members', []), - temporary=body.get('temporary'), - ) - await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce}) - - elif action == "get-chats": - chats_data = await app.get_chats(user_id) - await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce}) - - elif action == "delete-chat": - chat_id = await app.delete_chat( - chat_id=body['chatId'] - ) - await send({'message': "chat deleted", 'chatId': chat_id, 'nonce': nonce}) - - elif action == "get-chat-messages": - messages_data = await app.get_chat_messages( - body['chatId'], - earliest_index=body.get('earliestIndex', 0) - ) - await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) - - elif action == "get-chat-threads": - await send({'message': "threads retrieved", 'nonce': nonce, - 'threads': await app.get_chat_threads(body['chatId']) - }) - - elif action == "get-chat-thread-heads": - messages_data = await app.get_chat_thread_heads( - body['chatId'], - limit_message_id=body.get('limitMessageId', None) - ) - await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) - - elif action == "get-thread-messages": - messages_data = await app.get_thread_messages( - body['chatId'], - body['threadId'], - limit_message_id=body.get('limitMessageId', None) - ) - await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) - - elif action == "send-chat-message": - message = await app.send_chat_message( - chat_id=body['chatId'], - thread_id=body.get('threadId'), - content=body['content'], - kind=body.get('kind'), - subkind=body.get('subkind'), - format=body.get('format'), - ) - await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce}) - - elif action == "update-chat-message": - message = await app.update_chat_message( - message_id=body['messageId'], - content=body.get('content'), - kind=body.get('kind'), - subkind=body.get('subkind'), - format=body.get('format'), - ) - await send({'message': "message sent", 'messageId': body['messageId'], 'nonce': nonce}) - - elif action == "create-chat-thread": - thread_id = await app.create_chat_thread( - chat_id=body['chatId'], - name=body.get('name'), - ) - await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce}) - - elif action == "fork-chat-thread": - thread_id = await app.fork_chat_thread( - thread_id=body['threadId'], - message_id=body.get('messageId'), - ) - await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce}) - - elif action == "prompt-agent": - if 'instruction' not in body and 'input' not in body: - await send({'error': "instruction or input required", 'nonce': nonce}) - continue - - # Multiple-result action, same shape as read-stream: every - # generated message is its own packet carrying the request - # nonce, so the client callback runs once per message. The - # terminal eof packet resolves the action. Streams referenced - # by open assistant/reasoning messages are read by the - # client via the existing read-stream action. - generator = app.prompt_agent( - await agent_id_from_body(), - body.get('instruction'), - body.get('input'), - body.get('options', {}) - ) - try: - async for msg in generator: - await send({'message': 'agent-message', 'data': msg, 'nonce': nonce}) - finally: - # Disconnect / cancellation point: closes the run task. - await generator.aclose() - await send({'eof': True, 'nonce': nonce}) - - elif action == "read-stream": - stream_id = body['streamId'] - print("Read Stream" + str(stream_id)) - stream = app.streams.get(stream_id) - if stream is None: - await send({'error': 'bad stream id: ' + str(stream_id), 'nonce': nonce}) - else: - async def chunk_fn(chunk): - await send({'data': chunk, 'eof': False, 'nonce': nonce}) - async def eof_fn(): - await send({'eof': True, 'fullData': stream.content, 'nonce': nonce}) - stream.listen(chunk_fn, eof_fn) - - else: - await send({'error': 'unknown action: ' + action, 'nonce': nonce}) finally: print("User logged out:", user_id) sender_task.cancel() @@ -351,4 +86,294 @@ def register(server, app, auth_manager): # disconnected manually (possibly from another session) are safe. for remotename in list(session_connections): app.disconnectRemote(user_id, remotename) - app.remove_user_connection(user_id, send_queue) \ No newline at end of file + app.remove_user_connection(user_id, send_queue) + +# Error handling +async def handle_request(app, user_id, session_connections, send, req): + try: + req = json.loads(req) + await process_request(app, user_id, session_connections, send, req) + + except Exception as err: + action = req.get('action') if isinstance(req, dict) else None + nonce = req.get('nonce') if isinstance(req, dict) else None + + # Maintainers + print(f"[ws] req error @{user_id} {action!r}: {err!r}") + print(traceback.format_exc()) + # Client — nonce makes this routable to the caller's promise/callback + try: + await send({'error': str(err), 'nonce': nonce}) + except Exception: + print(traceback.format_exc()) + +# Actual processing +async def process_request(app, user_id, session_connections, send, req): + if 'action' not in req: + await send({'error': 'no action'}) + return + action = req['action'] + body = req.get('body', {}) + nonce = req.get('nonce', None) + print(f"[ws] req @{user_id} {action}") + + async def agent_id_from_body (allow_none=False): + if 'agentId' in body: + return body['agentId'] + elif 'agentname' in body: + return await app.agent_id_from_name(body['agentname']) + elif allow_none: + return None + else: + raise Exception("field 'agentId' or 'agentname' required") + + if action == "ping": + await send({'message': "pong", 'nonce': nonce}) + + elif action == "create-agent": + name = body['name'] + agent_id = await app.create_agent( + user_id=user_id, + name=name, + agentname=body.get('agentname'), + description=body.get('description'), + picture=body.get('picture'), + kind=body['kind'], + data=body.get('data'), + temporary=body.get('temporary'), + ) + await send({'message': "agent created", 'agentId': agent_id, 'nonce': nonce}) + + elif action == "update-agent": + agent_id = await agent_id_from_body() + await app.update_agent( + agent_id=agent_id, + name=body.get('name'), + agentname=body.get('agentname'), + description=body.get('description'), + picture=body.get('picture'), + data=body.get('data', {}) + ) + await send({'message': "agent updated", 'agentId': agent_id, 'nonce': nonce}) + + elif action == "get-models": + await send({ + 'models': [{"id": model.id, "name": model.name, "value": model.value, "kind": model.kind} for model in await app.get_models()], + 'nonce': nonce + }) + + elif action == "save-model": + model = await app.save_model( + model_id=body.get("modelId"), + name=body.get("name"), + kind=body.get("kind"), + value=body.get("value"), + ) + if model: + await send({'message': "model saved", 'modelId': model.id, 'nonce': nonce}) + elif "modelId" in body: + await send({'message': "model deleted", 'modelId': body["modelId"], 'nonce': nonce}) + + elif action == "get-agent-config-fields": + await send({'fields': await app.get_agent_config_fields(await agent_id_from_body()), 'nonce': nonce}) + + elif action == "get-agent-kinds": + await send({'message': "agents retrieved", 'kinds': app.get_agent_kinds(), 'nonce': nonce}) + + elif action == "get-agent-kind-info": + await send({'message': "agent kind retrieved", 'info': app.get_agent_kind_info(body['kind']), 'nonce': nonce}) + + elif action == "get-agents": + await send({'message': "agents retrieved", 'agents': await app.get_agents(user_id), 'nonce': nonce}) + + elif action == "get-agent-resources": + values = await app.get_agent_resources(await agent_id_from_body(), body.get('kind'), body.get('key')) + await send({'resources': [{'kind': r.kind, 'key': r.key, 'value': r.value} for r in values], 'nonce': nonce}) + + elif action == "save-agent-resource": + await app.save_agent_resource(await agent_id_from_body(), body['kind'], body['key'], body['value']) + await send({'message': "resource saved", 'nonce': nonce}) + + elif action == "delete-agent-resource": + await app.delete_agent_resource(await agent_id_from_body(), body['kind'], body['key']) + await send({'message': "resource deleted", 'nonce': nonce}) + + elif action == "get-agent-runs": + await send({ + 'runs': await app.get_agent_runs(agent_id=await agent_id_from_body(allow_none=True), user_id=user_id), + 'nonce': nonce + }) + + elif action == "cancel-agent-run": + await app.cancel_agent_run(body['runId']) + + elif action == "remote-list": + await send({'remotes': app.listRemotes(user_id), 'nonce': nonce}) + + elif action == "remote-connect": + # Connect a remote served by THIS session's channel. + # Expected failures (duplicate name) are replied, not + # raised, so the session stays alive. + try: + connection = app.connectRemote(user_id, body['remotename'], body.get('functions', []), send) + session_connections[body['remotename']] = connection + await send({'message': "remote connected", 'remotename': body['remotename'], 'nonce': nonce}) + except Exception as e: + await send({'error': str(e), 'nonce': nonce}) + + elif action == "remote-disconnect": + # Manual disconnect, allowed from any session of the user + connection = app.disconnectRemote(user_id, body['remotename']) + session_connections.pop(body['remotename'], None) + if connection is None: + await send({'error': "no connected remote named '" + str(body['remotename']) + "'", 'nonce': nonce}) + else: + await send({'message': "remote disconnected", 'remotename': body['remotename'], 'nonce': nonce}) + + elif action == "remote-call": + # Invoke a function of one of the user's connected remotes. + # Expected failures (unknown remote/function, timeout, + # remote-side error) are replied, not raised. + + + async def _handle_remote_call(user_id, remotename, func_name, args, send, nonce): + try: + result = await app.callRemote(user_id, remotename, func_name, args) + await send({'result': result, 'nonce': nonce}) + except Exception as e: + await send({'error': str(e), 'nonce': nonce}) + asyncio.create_task(_handle_remote_call( + user_id, body['remotename'], body['function'], body.get('args', {}), send, nonce + )) + + elif action == "remote-call-result": + # A session serving remotes replies to a downstream call. + # Fields live at the message top level, not in 'body'. + # requestIds are '#' and resolve/reject + # are no-ops for unknown ids, so fanning out over this + # session's connections routes the reply correctly. + app.resolve_remote_call(body['requestId'], body.get('result'), error=body.get('error')) + + elif action == "create-chat": + chat_id = await app.create_chat( + user_id=user_id, + name=body['name'], + description=body.get('description'), + picture=body.get('picture'), + members=body.get('members', []), + temporary=body.get('temporary'), + ) + await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce}) + + elif action == "get-chats": + chats_data = await app.get_chats(user_id) + await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce}) + + elif action == "delete-chat": + chat_id = await app.delete_chat( + chat_id=body['chatId'] + ) + await send({'message': "chat deleted", 'chatId': chat_id, 'nonce': nonce}) + + elif action == "get-chat-messages": + messages_data = await app.get_chat_messages( + body['chatId'], + earliest_index=body.get('earliestIndex', 0) + ) + await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) + + elif action == "get-chat-threads": + await send({'message': "threads retrieved", 'nonce': nonce, + 'threads': await app.get_chat_threads(body['chatId']) + }) + + elif action == "get-chat-thread-heads": + messages_data = await app.get_chat_thread_heads( + body['chatId'], + limit_message_id=body.get('limitMessageId', None) + ) + await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) + + elif action == "get-thread-messages": + messages_data = await app.get_thread_messages( + body['chatId'], + body['threadId'], + limit_message_id=body.get('limitMessageId', None) + ) + await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce}) + + elif action == "send-chat-message": + message = await app.send_chat_message( + chat_id=body['chatId'], + thread_id=body.get('threadId'), + content=body['content'], + kind=body.get('kind'), + subkind=body.get('subkind'), + format=body.get('format'), + ) + await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce}) + + elif action == "update-chat-message": + message = await app.update_chat_message( + message_id=body['messageId'], + content=body.get('content'), + kind=body.get('kind'), + subkind=body.get('subkind'), + format=body.get('format'), + ) + await send({'message': "message sent", 'messageId': body['messageId'], 'nonce': nonce}) + + elif action == "create-chat-thread": + thread_id = await app.create_chat_thread( + chat_id=body['chatId'], + name=body.get('name'), + ) + await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce}) + + elif action == "fork-chat-thread": + thread_id = await app.fork_chat_thread( + thread_id=body['threadId'], + message_id=body.get('messageId'), + ) + await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce}) + + elif action == "prompt-agent": + if 'instruction' not in body and 'input' not in body: + await send({'error': "instruction or input required", 'nonce': nonce}) + return + + # Multiple-result action, same shape as read-stream: every + # generated message is its own packet carrying the request + # nonce, so the client callback runs once per message. The + # terminal eof packet resolves the action. Streams referenced + # by open assistant/reasoning messages are read by the + # client via the existing read-stream action. + generator = app.prompt_agent( + await agent_id_from_body(), + body.get('instruction'), + body.get('input'), + body.get('options', {}) + ) + try: + async for msg in generator: + await send({'message': 'agent-message', 'data': msg, 'nonce': nonce}) + finally: + # Disconnect / cancellation point: closes the run task. + await generator.aclose() + await send({'eof': True, 'nonce': nonce}) + + elif action == "read-stream": + stream_id = body['streamId'] + print("Read Stream" + str(stream_id)) + stream = app.streams.get(stream_id) + if stream is None: + await send({'error': 'bad stream id: ' + str(stream_id), 'nonce': nonce}) + else: + async def chunk_fn(chunk): + await send({'data': chunk, 'eof': False, 'nonce': nonce}) + async def eof_fn(): + await send({'eof': True, 'fullData': stream.content, 'nonce': nonce}) + stream.listen(chunk_fn, eof_fn) + + else: + await send({'error': 'unknown action: ' + action, 'nonce': nonce}) \ No newline at end of file