diff --git a/src/api-service/__app__/node/__init__.py b/src/api-service/__app__/node/__init__.py index 78a9310c6..6e122e7cd 100644 --- a/src/api-service/__app__/node/__init__.py +++ b/src/api-service/__app__/node/__init__.py @@ -12,7 +12,7 @@ from onefuzztypes.responses import BoolResult from ..onefuzzlib.endpoint_authorization import call_if_user from ..onefuzzlib.events import get_events from ..onefuzzlib.request import not_ok, ok, parse_request -from ..onefuzzlib.workers.nodes import Node, NodeTasks +from ..onefuzzlib.workers.nodes import Node, NodeMessage, NodeTasks def get(req: func.HttpRequest) -> func.HttpResponse: @@ -33,6 +33,9 @@ def get(req: func.HttpRequest) -> func.HttpResponse: node_tasks = NodeTasks.get_by_machine_id(request.machine_id) node.tasks = [(t.task_id, t.state) for t in node_tasks] + node.messages = [ + x.message for x in NodeMessage.get_messages(request.machine_id) + ] return ok(node) diff --git a/src/api-service/__app__/onefuzzlib/workers/nodes.py b/src/api-service/__app__/onefuzzlib/workers/nodes.py index a3239e265..7858688f5 100644 --- a/src/api-service/__app__/onefuzzlib/workers/nodes.py +++ b/src/api-service/__app__/onefuzzlib/workers/nodes.py @@ -157,7 +157,7 @@ class Node(BASE_NODE, ORMMixin): return ("pool_name", "machine_id") def save_exclude(self) -> Optional[MappingIntStrAny]: - return {"tasks": ...} + return {"tasks": ..., "messages": ...} def telemetry_include(self) -> Optional[MappingIntStrAny]: return { diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py index f701a6a9a..2489abc3a 100644 --- a/src/pytypes/onefuzztypes/models.py +++ b/src/pytypes/onefuzztypes/models.py @@ -557,6 +557,29 @@ class NodeHeartbeatEntry(BaseModel): data: List[Dict[str, HeartbeatType]] +class StopNodeCommand(BaseModel): + pass + + +class StopTaskNodeCommand(BaseModel): + task_id: UUID + + +class NodeCommandAddSshKey(BaseModel): + public_key: str + + +class NodeCommand(EnumModel): + stop: Optional[StopNodeCommand] + stop_task: Optional[StopTaskNodeCommand] + add_ssh_key: Optional[NodeCommandAddSshKey] + + +class NodeCommandEnvelope(BaseModel): + command: NodeCommand + message_id: str + + class Node(BaseModel): timestamp: Optional[datetime] = Field(alias="Timestamp") pool_name: PoolName @@ -564,6 +587,7 @@ class Node(BaseModel): state: NodeState = Field(default=NodeState.init) scaleset_id: Optional[UUID] = None tasks: Optional[List[Tuple[UUID, NodeTaskState]]] = None + messages: Optional[List[NodeCommand]] = None heartbeat: Optional[datetime] version: str = Field(default="1.0.0") reimage_requested: bool = Field(default=False) @@ -776,29 +800,6 @@ class NodeEventEnvelope(BaseModel): event: NodeEventShim -class StopNodeCommand(BaseModel): - pass - - -class StopTaskNodeCommand(BaseModel): - task_id: UUID - - -class NodeCommandAddSshKey(BaseModel): - public_key: str - - -class NodeCommand(EnumModel): - stop: Optional[StopNodeCommand] - stop_task: Optional[StopTaskNodeCommand] - add_ssh_key: Optional[NodeCommandAddSshKey] - - -class NodeCommandEnvelope(BaseModel): - command: NodeCommand - message_id: str - - class TaskEvent(BaseModel): timestamp: Optional[datetime] = Field(alias="Timestamp") task_id: UUID