Coverage for core / src / sensorkit / core / impl / device.py: 88%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5from _contextvars import ContextVar 

6from typing import Callable, ClassVar, override 

7 

8from loguru import logger 

9from pydantic import ValidationError 

10 

11from sensorkit.backend.request import CallContext 

12from sensorkit.core.device import ( 

13 CommandDone, 

14 CommandHandlerCallback, 

15 CommandRequestMessage, 

16 CommandResult, 

17 CommandStarted, 

18 DeviceCommand, 

19 DeviceEnableState, 

20 DeviceEnableStateRequest, 

21 DeviceInterface, 

22 DeviceState, 

23 run_command_request, 

24 set_enable_state_request, 

25) 

26from sensorkit.core.entity import DeviceDetails, EntityInfo 

27from sensorkit.core.impl.entity import EntityImpl 

28 

29 

30class DeviceImpl(EntityImpl, DeviceInterface): 

31 """Helper for implementing server-side functionality of a Device.""" 

32 

33 current: ClassVar[ContextVar[DeviceImpl | None]] = ContextVar("current_device", default=None) 

34 

35 def __init__(self, **kwargs): 

36 super().__init__(**kwargs) 

37 

38 self._enable_hooks: list[Callable[[], None]] = [] 

39 self._disable_hooks: list[Callable[[], None]] = [] 

40 self._handlers: dict[str, CommandHandlerCallback] = {} 

41 self._published_keywords: set[str] = set() 

42 

43 def declare_published_keyword(self, keyword_id: str): 

44 """Declare that this device publishes a keyword.""" 

45 self._published_keywords.add(keyword_id) 

46 

47 @override 

48 def on_enable(self, func: Callable[[], None]): 

49 self._enable_hooks.append(func) 

50 return func 

51 

52 @override 

53 def on_disable(self, func: Callable[[], None]): 

54 self._disable_hooks.append(func) 

55 return func 

56 

57 @override 

58 async def init_impl(self): 

59 self._state = await DeviceState.recover_or_init( 

60 self, 

61 enable_state=DeviceEnableState(enabled=True), 

62 ) 

63 

64 @override 

65 async def attach_impl(self): 

66 if self._state.enable_state.enabled: 

67 await self._call_with_context(self._enable_hooks) 

68 

69 await self.handle_request(set_enable_state_request, self._set_enable_state) 

70 await self.handle_request(run_command_request, self._command_request) 

71 

72 async def _set_enable_state(self, request: DeviceEnableStateRequest): 

73 if request.enable == self._state.enable_state.enabled: 

74 return 

75 

76 await self._state.update(self, DeviceEnableState(enabled=request.enable)) 

77 

78 if request.enable: 

79 await self._call_with_context(self._enable_hooks) 

80 else: 

81 # TODO: Interrupt all ongoing commands. 

82 

83 await self._call_with_context(self._disable_hooks) 

84 

85 async def _command_request( 

86 self, 

87 message: CommandRequestMessage, 

88 call: CallContext[None, CommandResult], 

89 ): 

90 command_id = message.command.command_id 

91 

92 # Reject if we aren't ready. 

93 if not self._state.enable_state.enabled: 

94 logger.warning(f"Rejecting {command_id} command: Device is disabled") 

95 call.reject(response=None) 

96 return 

97 

98 # Look for a handler for this command ID. 

99 if command_id not in self._handlers: 

100 logger.warning(f"Rejecting unhandled command: {command_id}") 

101 call.reject(response=None) 

102 return 

103 

104 # Accept the command and invoke the configured handler func. 

105 call.accept(response=None) 

106 handler_func = self._handlers[command_id] 

107 success = False 

108 

109 # Emit the command start event. 

110 logger.debug(f"Incoming {command_id} Command: {call.call_id}") 

111 await self.emit_event(CommandStarted(command_id=command_id, call_id=call.call_id)) 

112 

113 try: 

114 # Invoke the registered handler function in a new task and store a reference. 

115 with self.enter_context(): 

116 task = asyncio.create_task(handler_func(message.command)) 

117 

118 # Wait for the handler to return. 

119 await call.progress_from_task(task, cadence=6.0, ttl=10.0) 

120 success = True 

121 except ValidationError: 

122 logger.exception(f"Incoming {command_id} Command failed validation") 

123 raise 

124 except asyncio.CancelledError: 

125 logger.exception(f"Execution of {command_id} Command cancelled") 

126 raise 

127 except Exception: 

128 logger.exception(f"Error invoking {command_id} Command handler") 

129 raise 

130 else: 

131 # The command completed successfully, so return the result including the handler func 

132 # return value. 

133 await call.succeed(result=CommandResult(data=task.result())) 

134 finally: 

135 # Emit the command end event. 

136 try: 

137 await self.emit_event( 

138 CommandDone( 

139 command_id=command_id, 

140 call_id=call.call_id, 

141 success=success, 

142 ) 

143 ) 

144 except Exception as e: 

145 logger.warning(f"Error logging command done event ({type(e).__name__})") 

146 

147 @override 

148 def command_handler(self, command_type: type[DeviceCommand]): 

149 key = command_type.model_tag() 

150 

151 def decorator(func: CommandHandlerCallback): 

152 self._handlers[key] = func 

153 return func 

154 

155 return decorator 

156 

157 @override 

158 def entity_info(self) -> EntityInfo: 

159 return EntityInfo( 

160 entity_type="device", 

161 details=DeviceDetails( 

162 supported_commands=frozenset(self._handlers.keys()), 

163 published_keywords=frozenset(self._published_keywords), 

164 ), 

165 )