def __init__(self, server):
self._server = server
self._shutdown = False
self._client_request_futures = {}
self._server_request_futures = {}
self.fm = FeatureManager(server)
self.transport = None
self._message_buf = []
class JsonRPCProtocol(asyncio.Protocol):
"""Json RPC protocol implementation using on top of `asyncio.Protocol`.
Specification of the protocol can be found here:
https://www.jsonrpc.org/specification
This class provides bidirectional communication which is needed for LSP.
"""
CHARSET = 'utf-8'
CONTENT_TYPE = 'application/vscode-jsonrpc'
MESSAGE_PATTERN = re.compile(
rb'^(?:[^\r\n]+\r\n)*'
+ rb'Content-Length: (?P<length>\d+)\r\n'
+ rb'(?:[^\r\n]+\r\n)*\r\n'
+ rb'(?P<body>{.*)',
re.DOTALL,
)
VERSION = '2.0'
def __init__(self, server):
self._server = server
self._shutdown = False
self._client_request_futures = {}
self._server_request_futures = {}
self.fm = FeatureManager(server)
self.transport = None
self._message_buf = []
self._send_only_body = False
def __call__(self):
return self
def _check_ret_type_and_send_response(self, method_name, method_type, msg_id, result):
"""Check if registered feature returns appropriate result type."""
if method_type == ATTR_FEATURE_TYPE:
return_type = get_method_return_type(method_name)
if not is_instance(result, return_type):
error = JsonRpcInternalError().to_dict()
self._send_response(msg_id, error=error)
self._send_response(msg_id, result=result)
def _execute_notification(self, handler, *params):
"""Executes notification message handler."""
if asyncio.iscoroutinefunction(handler):
future = asyncio.ensure_future(handler(*params))
future.add_done_callback(self._execute_notification_callback)
else:
if is_thread_function(handler):
self._server.thread_pool.apply_async(handler, (*params, ))
else:
handler(*params)
def _execute_notification_callback(self, future):
"""Success callback used for coroutine notification message."""
if future.exception():
try:
raise future.exception()
except Exception:
error = JsonRpcInternalError.of(sys.exc_info()).to_dict()
logger.exception('Exception occurred in notification: "%s"', error)
# Revisit. Client does not support response with msg_id = None
# https://stackoverflow.com/questions/31091376/json-rpc-2-0-allow-notifications-to-have-an-error-response
# self._send_response(None, error=error)
def _execute_request(self, msg_id, handler, params):
"""Executes request message handler."""
method_name, method_type = get_help_attrs(handler)
if asyncio.iscoroutinefunction(handler):
future = asyncio.ensure_future(handler(params))
self._client_request_futures[msg_id] = future
future.add_done_callback(partial(self._execute_request_callback,
method_name, method_type, msg_id))
else:
# Can't be canceled
if is_thread_function(handler):
self._server.thread_pool.apply_async(
handler, (params, ),
callback=partial(
self._check_ret_type_and_send_response,
method_name, method_type, msg_id,
),
error_callback=partial(self._execute_request_err_callback, msg_id))
else:
self._check_ret_type_and_send_response(
method_name, method_type, msg_id, handler(params))
def _execute_request_callback(self, method_name, method_type, msg_id, future):
"""Success callback used for coroutine request message."""
try:
if not future.cancelled():
self._check_ret_type_and_send_response(
method_name, method_type, msg_id, result=future.result())
else:
self._send_response(
msg_id,
error=JsonRpcRequestCancelled(f'Request with id "{msg_id}" is canceled')
)
self._client_request_futures.pop(msg_id, None)
except Exception:
error = JsonRpcInternalError.of(sys.exc_info()).to_dict()
logger.exception('Exception occurred for message "%s": %s', msg_id, error)
self._send_response(msg_id, error=error)
def _execute_request_err_callback(self, msg_id, exc):
"""Error callback used for coroutine request message."""
exc_info = (type(exc), exc, None)
error = JsonRpcInternalError.of(exc_info).to_dict()
logger.exception('Exception occurred for message "%s": %s', msg_id, error)
self._send_response(msg_id, error=error)
def _get_handler(self, feature_name):
"""Returns builtin or used defined feature by name if exists."""
try:
return self.fm.builtin_features[feature_name]
except KeyError:
try:
return self.fm.features[feature_name]
except KeyError:
raise JsonRpcMethodNotFound.of(feature_name)
def _handle_cancel_notification(self, msg_id):
"""Handles a cancel notification from the client."""
future = self._client_request_futures.pop(msg_id, None)
if not future:
logger.warning('Cancel notification for unknown message id "%s"', msg_id)
return
# Will only work if the request hasn't started executing
if future.cancel():
logger.info('Cancelled request with id "%s"', msg_id)
def _handle_notification(self, method_name, params):
"""Handles a notification from the client."""
if method_name == CANCEL_REQUEST:
self._handle_cancel_notification(params.id)
return
try:
handler = self._get_handler(method_name)
self._execute_notification(handler, params)
except (KeyError, JsonRpcMethodNotFound):
logger.warning('Ignoring notification for unknown method "%s"', method_name)
except Exception:
logger.exception('Failed to handle notification "%s": %s', method_name, params)
def _handle_request(self, msg_id, method_name, params):
"""Handles a request from the client."""
try:
handler = self._get_handler(method_name)
# workspace/executeCommand is a special case
if method_name == WORKSPACE_EXECUTE_COMMAND:
handler(params, msg_id)
else:
self._execute_request(msg_id, handler, params)
except JsonRpcException as e:
logger.exception('Failed to handle request %s %s %s', msg_id, method_name, params)
self._send_response(msg_id, None, e.to_dict())
except Exception:
logger.exception('Failed to handle request %s %s %s', msg_id, method_name, params)
err = JsonRpcInternalError.of(sys.exc_info()).to_dict()
self._send_response(msg_id, None, err)
def _handle_response(self, msg_id, result=None, error=None):
"""Handles a response from the client."""
future = self._server_request_futures.pop(msg_id, None)
if not future:
logger.warning('Received response to unknown message id "%s"', msg_id)
return
if error is not None:
logger.debug('Received error response to message "%s": %s', msg_id, error)
future.set_exception(JsonRpcException.from_dict(error))
else:
logger.debug('Received result for message "%s": %s', msg_id, result)
future.set_result(result)
def _procedure_handler(self, message):
"""Delegates message to handlers depending on message type."""
if message.jsonrpc != JsonRPCProtocol.VERSION:
logger.warning('Unknown message "%s"', message)
return
if self._shutdown and getattr(message, 'method', '') != EXIT:
logger.warning('Server shutting down. No more requests!')
return
if isinstance(message, JsonRPCNotification):
logger.debug('Notification message received.')
self._handle_notification(message.method, message.params)
elif isinstance(message, JsonRPCResponseMessage):
logger.debug('Response message received.')
self._handle_response(message.id, message.result, message.error)
elif isinstance(message, JsonRPCRequestMessage):
logger.debug('Request message received.')
self._handle_request(message.id, message.method, message.params)
def _send_data(self, data):
"""Sends data to the client."""
if not data:
return
try:
body = data.json(by_alias=True, exclude_unset=True, encoder=default_serializer)
logger.info('Sending data: %s', body)
body = body.encode(self.CHARSET)
if not self._send_only_body:
header = (
f'Content-Length: {len(body)}\r\n'
f'Content-Type: {self.CONTENT_TYPE}; charset={self.CHARSET}\r\n\r\n'
).encode(self.CHARSET)
self.transport.write(header + body)
else:
self.transport.write(body.decode('utf-8'))
except Exception:
logger.error(traceback.format_exc())
def _send_response(self, msg_id, result=None, error=None):
"""Sends a JSON RPC response to the client.
Args:
msg_id(str): Id from request
result(any): Result returned by handler
error(any): Error returned by handler
"""
response = JsonRPCResponseMessage(id=msg_id,
jsonrpc=JsonRPCProtocol.VERSION,
result=result,
error=error)
if error is None:
del response.error
else:
del response.result
self._send_data(response)
def connection_lost(self, exc):
"""Method from base class, called when connection is lost, in which case we
want to shutdown the server's process as well.
"""
logger.error('Connection to the client is lost! Shutting down the server.')
sys.exit(1)
def connection_made(self, transport: asyncio.BaseTransport):
"""Method from base class, called when connection is established"""
self.transport = transport
def data_received(self, data: bytes):
"""Method from base class, called when server receives the data"""
logger.debug('Received %r', data)
while len(data):
# Append the incoming chunk to the message buffer
self._message_buf.append(data)
# Look for the body of the message
message = b''.join(self._message_buf)
found = JsonRPCProtocol.MESSAGE_PATTERN.fullmatch(message)
body = found.group('body') if found else b''
length = int(found.group('length')) if found else 1
if len(body) < length:
# Message is incomplete; bail until more data arrives
return
# Message is complete;
# extract the body and any remaining data,
# and reset the buffer for the next message
body, data = body[:length], body[length:]
self._message_buf = []
# Parse the body
self._procedure_handler(
json.loads(body.decode(self.CHARSET),
object_hook=deserialize_message))
def notify(self, method: str, params=None):
"""Sends a JSON RPC notification to the client."""
logger.debug('Sending notification: "%s" %s', method, params)
request = JsonRPCNotification(
jsonrpc=JsonRPCProtocol.VERSION,
method=method,
params=params
)
self._send_data(request)
def send_request(self, method, params=None, callback=None):
"""Sends a JSON RPC request to the client.
Args:
method(str): The method name of the message to send
params(any): The payload of the message
Returns:
Future that will be resolved once a response has been received
"""
msg_id = str(uuid.uuid4())
logger.debug('Sending request with id "%s": %s %s', msg_id, method, params)
request = JsonRPCRequestMessage(
id=msg_id,
jsonrpc=JsonRPCProtocol.VERSION,
method=method,
params=params
)
future = Future()
# If callback function is given, call it when result is received
if callback:
def wrapper(future: Future):
result = future.result()
logger.info('Client response for %s received: %s', params, result)
callback(result)
future.add_done_callback(wrapper)
self._server_request_futures[msg_id] = future
self._send_data(request)
return future
def send_request_async(self, method, params=None):
"""Calls `send_request` and wraps `concurrent.futures.Future` with
`asyncio.Future` so it can be used with `await` keyword.
Args:
method(str): The method name of the message to send
params(any): The payload of the message
Returns:
`asyncio.Future` that can be awaited
"""
return asyncio.wrap_future(self.send_request(method, params))
def thread(self):
"""Decorator that mark function to execute it in a thread."""
return self.fm.thread()
def feature_manager():
""" Return a feature manager """
return FeatureManager()