implement heartbeat command

This commit is contained in:
ars
2026-05-24 18:36:55 +03:00
parent c92c5195ce
commit f9aed123fd
11 changed files with 663 additions and 25 deletions
+1
View File
@@ -19,6 +19,7 @@ dependencies = [
[project.optional-dependencies] [project.optional-dependencies]
dev = [ dev = [
"pytest>=8.0", "pytest>=8.0",
"httpx>=0.28.0",
"black>=24.0", "black>=24.0",
"pyinstaller>=6.0", "pyinstaller>=6.0",
"mypy>=1.0", "mypy>=1.0",
+127 -16
View File
@@ -1,23 +1,134 @@
import shlex import platform
import subprocess import subprocess
from typing import Dict, Optional from datetime import UTC, datetime
from typing import Any, Callable, Mapping, Sequence
from nexus_sync.common import Command, CommandKind, CommandResult, CommandResultStatus
PresetBuilder = Callable[[Mapping[str, Any]], Sequence[str]]
DEFAULT_OUTPUT_LIMIT_BYTES = 64 * 1024
def run(command: str, stdin: Optional[str] = None) -> Dict[str, str]: def _reject(command: Command, message: str) -> CommandResult:
process = _execute(command, stdin) now = datetime.now(UTC)
return CommandResult(
return { command_id=command.id,
"stdout": process.stdout, status=CommandResultStatus.REJECTED,
"stderr": process.stderr, started_at=now,
} finished_at=now,
stderr=message,
)
def _parse_command(command: str) -> list[str]: def _network_interfaces(args: Mapping[str, Any]) -> Sequence[str]:
# return command.split() if args:
return shlex.split(command) raise ValueError("network_interfaces does not accept arguments")
system = platform.system().lower()
if system == "windows":
return ["ipconfig"]
if system == "linux":
return ["ip", "addr"]
if system == "darwin":
return ["ifconfig"]
raise ValueError(f"unsupported platform for network_interfaces: {system or 'unknown'}")
def _execute(command: str, stdin: Optional[str] = None) -> subprocess.CompletedProcess: def _hostname(args: Mapping[str, Any]) -> Sequence[str]:
parsed_command = _parse_command(command) if args:
result = subprocess.run(parsed_command, input=stdin, text=True, capture_output=True, shell=True) raise ValueError("hostname does not accept arguments")
return result return ["hostname"]
DEFAULT_PRESETS: dict[str, PresetBuilder] = {
"hostname": _hostname,
"network_interfaces": _network_interfaces,
}
def execute_command(
command: Command,
*,
stdin: str | None = None,
presets: Mapping[str, PresetBuilder] = DEFAULT_PRESETS,
output_limit_bytes: int = DEFAULT_OUTPUT_LIMIT_BYTES,
) -> CommandResult:
if command.kind != CommandKind.EXEC:
return _reject(command, f"unsupported command kind: {command.kind}")
builder = presets.get(command.name)
if builder is None:
return _reject(command, f"unknown command preset: {command.name}")
try:
argv = list(builder(command.args))
except ValueError as error:
return _reject(command, str(error))
if not argv:
return _reject(command, f"command preset returned empty argv: {command.name}")
started_at = datetime.now(UTC)
try:
process = subprocess.run(
argv,
input=stdin,
text=True,
capture_output=True,
timeout=command.timeout_seconds,
shell=False,
check=False,
)
except subprocess.TimeoutExpired as error:
finished_at = datetime.now(UTC)
return CommandResult(
command_id=command.id,
status=CommandResultStatus.TIMED_OUT,
started_at=started_at,
finished_at=finished_at,
stdout=_limit_output(error.stdout or "", output_limit_bytes),
stderr=_limit_output(error.stderr or "", output_limit_bytes),
)
except OSError as error:
finished_at = datetime.now(UTC)
return CommandResult(
command_id=command.id,
status=CommandResultStatus.FAILED,
started_at=started_at,
finished_at=finished_at,
return_code=None,
stderr=str(error),
)
finished_at = datetime.now(UTC)
status = (
CommandResultStatus.SUCCEEDED if process.returncode == 0 else CommandResultStatus.FAILED
)
return CommandResult(
command_id=command.id,
status=status,
started_at=started_at,
finished_at=finished_at,
return_code=process.returncode,
stdout=_limit_output(process.stdout, output_limit_bytes),
stderr=_limit_output(process.stderr, output_limit_bytes),
)
def _limit_output(value: str | bytes, limit_bytes: int) -> str:
if isinstance(value, bytes):
value = value.decode(errors="replace")
encoded = value.encode()
if len(encoded) <= limit_bytes:
return value
marker = "\n[truncated]"
marker_bytes = marker.encode()
if limit_bytes <= len(marker_bytes):
return encoded[:limit_bytes].decode(errors="ignore")
content_limit = max(0, limit_bytes - len(marker_bytes))
truncated = encoded[:content_limit].decode(errors="ignore")
return f"{truncated}{marker}"
+11
View File
@@ -0,0 +1,11 @@
from nexus_sync.server.app import create_app
from nexus_sync.server.config import DEFAULT_COMMAND_POLL_SECONDS, DEFAULT_IDLE_POLL_SECONDS
from nexus_sync.server.store import InMemoryStore, Store
__all__ = [
"DEFAULT_COMMAND_POLL_SECONDS",
"DEFAULT_IDLE_POLL_SECONDS",
"InMemoryStore",
"Store",
"create_app",
]
+6 -1
View File
@@ -1,5 +1,10 @@
import uvicorn
from nexus_sync.server.app import app
def main() -> None: def main() -> None:
print("nexus-sync server") uvicorn.run(app, host="0.0.0.0", port=8000)
if __name__ == "__main__": if __name__ == "__main__":
+84
View File
@@ -0,0 +1,84 @@
from datetime import UTC, datetime
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from nexus_sync.common import HeartbeatRequest, HeartbeatResponse
from nexus_sync.server.config import (
DEFAULT_COMMAND_POLL_SECONDS,
DEFAULT_IDLE_POLL_SECONDS,
load_client_tokens,
)
from nexus_sync.server.store import InMemoryStore, Store
def create_app(
*,
store: Store | None = None,
client_tokens: dict[str, str] | None = None,
) -> FastAPI:
app = FastAPI(title="nexus-sync", version="0.1.0")
app.state.store = store or InMemoryStore()
app.state.client_tokens = client_tokens if client_tokens is not None else load_client_tokens()
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(
_request: Request,
error: RequestValidationError,
) -> JSONResponse:
return JSONResponse(
status_code=status.HTTP_400_BAD_REQUEST, content={"detail": error.errors()}
)
def authorize_token(authorization: str | None = Header(default=None)) -> str:
if authorization is None or not authorization.startswith("Bearer "):
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="missing bearer token",
)
token = authorization.removeprefix("Bearer ").strip()
owner_client_id = _client_id_for_token(app.state.client_tokens, token)
if owner_client_id is None:
raise HTTPException(
status_code=status.HTTP_401_UNAUTHORIZED,
detail="invalid bearer token",
)
return owner_client_id
@app.post("/api/v1/client/heartbeat", response_model=HeartbeatResponse)
def heartbeat(
payload: HeartbeatRequest,
token_client_id: str = Depends(authorize_token),
) -> HeartbeatResponse:
if token_client_id != payload.client_id:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN,
detail="token is not allowed for this client_id",
)
now = datetime.now(UTC)
app.state.store.upsert_client(payload, now)
app.state.store.record_command_result(payload, now)
command = app.state.store.take_next_command(payload.client_id, now)
return HeartbeatResponse(
server_time=now,
next_poll_after_seconds=(
DEFAULT_COMMAND_POLL_SECONDS if command else DEFAULT_IDLE_POLL_SECONDS
),
command=command,
)
return app
def _client_id_for_token(client_tokens: dict[str, str], token: str) -> str | None:
for client_id, expected_token in client_tokens.items():
if token == expected_token:
return client_id
return None
app = create_app()
+16
View File
@@ -0,0 +1,16 @@
import os
DEFAULT_IDLE_POLL_SECONDS = 60
DEFAULT_COMMAND_POLL_SECONDS = 10
def load_client_tokens() -> dict[str, str]:
raw_tokens = os.environ.get("NEXUS_SYNC_CLIENT_TOKENS", "dev-client:dev-token")
tokens: dict[str, str] = {}
for item in raw_tokens.split(","):
if not item.strip():
continue
client_id, separator, token = item.partition(":")
if separator and client_id and token:
tokens[client_id] = token
return tokens
+143
View File
@@ -0,0 +1,143 @@
from dataclasses import dataclass, field
from datetime import datetime
from typing import Protocol
from fastapi import HTTPException, status
from nexus_sync.common import (
ClientRecord,
Command,
CommandRecord,
CommandResultRecord,
CommandStatus,
HeartbeatRequest,
)
class Store(Protocol):
def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
pass
def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
pass
def take_next_command(self, client_id: str, now: datetime) -> Command | None:
pass
@dataclass
class InMemoryStore:
clients: dict[str, ClientRecord] = field(default_factory=dict)
commands: dict[str, CommandRecord] = field(default_factory=dict)
results: list[CommandResultRecord] = field(default_factory=list)
def upsert_client(self, heartbeat: HeartbeatRequest, now: datetime) -> ClientRecord:
existing = self.clients.get(heartbeat.client_id)
created_at = existing.created_at if existing else now
token_hash = existing.token_hash if existing else None
record = ClientRecord(
id=heartbeat.client_id,
hostname=heartbeat.client.hostname,
platform=heartbeat.client.platform,
version=heartbeat.client.version,
created_at=created_at,
last_seen_at=now,
is_active=True,
token_hash=token_hash,
)
self.clients[heartbeat.client_id] = record
return record
def enqueue_command(self, command: Command, client_id: str, now: datetime) -> CommandRecord:
record = CommandRecord(
id=command.id,
client_id=client_id,
kind=command.kind,
name=command.name,
args=command.args,
status=CommandStatus.PENDING,
timeout_seconds=command.timeout_seconds,
created_at=now,
)
self.commands[record.id] = record
return record
def record_command_result(self, heartbeat: HeartbeatRequest, now: datetime) -> None:
result = heartbeat.last_command_result
if result is None:
return
command = self.commands.get(result.command_id)
if command is None:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail="command result references unknown command",
)
if command.client_id != heartbeat.client_id:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail="command result references command for another client",
)
if command.status in TERMINAL_COMMAND_STATUSES:
raise HTTPException(
status_code=status.HTTP_409_CONFLICT,
detail="command result references terminal command",
)
finished_at = result.finished_at or now
self.commands[command.id] = command.model_copy(
update={
"status": CommandStatus(result.status.value),
"finished_at": finished_at,
}
)
self.results.append(
CommandResultRecord(
command_id=result.command_id,
client_id=heartbeat.client_id,
status=result.status,
started_at=result.started_at,
finished_at=finished_at,
return_code=result.return_code,
stdout=result.stdout,
stderr=result.stderr,
received_at=now,
)
)
def take_next_command(self, client_id: str, now: datetime) -> Command | None:
pending = sorted(
(
command
for command in self.commands.values()
if command.client_id == client_id and command.status == CommandStatus.PENDING
),
key=lambda command: command.created_at,
)
if not pending:
return None
record = pending[0]
self.commands[record.id] = record.model_copy(
update={
"status": CommandStatus.DELIVERED,
"attempts": record.attempts + 1,
"delivered_at": now,
}
)
return Command(
id=record.id,
kind=record.kind,
name=record.name,
args=record.args,
timeout_seconds=record.timeout_seconds,
)
TERMINAL_COMMAND_STATUSES = {
CommandStatus.SUCCEEDED,
CommandStatus.FAILED,
CommandStatus.TIMED_OUT,
CommandStatus.REJECTED,
}
+100
View File
@@ -0,0 +1,100 @@
import subprocess
from nexus_sync.client.execute import execute_command
from nexus_sync.common import Command, CommandKind, CommandResultStatus
def _command(name: str = "hostname", timeout_seconds: int = 30) -> Command:
return Command(
id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B",
kind=CommandKind.EXEC,
name=name,
args={},
timeout_seconds=timeout_seconds,
)
def test_execute_command_runs_known_preset_without_shell(monkeypatch) -> None:
calls = []
def fake_run(argv, **kwargs):
calls.append((argv, kwargs))
return subprocess.CompletedProcess(argv, 0, stdout="host\n", stderr="")
monkeypatch.setattr(subprocess, "run", fake_run)
result = execute_command(_command())
assert result.status == CommandResultStatus.SUCCEEDED
assert result.return_code == 0
assert result.stdout == "host\n"
assert calls == [
(
["hostname"],
{
"input": None,
"text": True,
"capture_output": True,
"timeout": 30,
"shell": False,
"check": False,
},
)
]
def test_execute_command_rejects_unknown_preset() -> None:
result = execute_command(_command(name="rm_everything"))
assert result.status == CommandResultStatus.REJECTED
assert result.return_code is None
assert "unknown command preset" in result.stderr
def test_execute_command_maps_non_zero_exit_to_failed(monkeypatch) -> None:
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(argv, 2, stdout="", stderr="failed\n")
monkeypatch.setattr(subprocess, "run", fake_run)
result = execute_command(_command())
assert result.status == CommandResultStatus.FAILED
assert result.return_code == 2
assert result.stderr == "failed\n"
def test_execute_command_maps_timeout(monkeypatch) -> None:
def fake_run(argv, **kwargs):
raise subprocess.TimeoutExpired(argv, timeout=kwargs["timeout"], output="partial")
monkeypatch.setattr(subprocess, "run", fake_run)
result = execute_command(_command(timeout_seconds=1))
assert result.status == CommandResultStatus.TIMED_OUT
assert result.return_code is None
assert result.started_at is not None
assert result.finished_at is not None
def test_execute_command_truncates_output(monkeypatch) -> None:
def fake_run(argv, **kwargs):
return subprocess.CompletedProcess(argv, 0, stdout="abcdefghijklmnopqrstuvwxyz", stderr="")
monkeypatch.setattr(subprocess, "run", fake_run)
result = execute_command(_command(), output_limit_bytes=20)
assert result.stdout.endswith("\n[truncated]")
assert len(result.stdout.encode()) <= 20
def test_execute_command_rejects_preset_validation_error() -> None:
command = _command()
invalid = command.model_copy(update={"args": {"unexpected": True}})
result = execute_command(invalid)
assert result.status == CommandResultStatus.REJECTED
assert "does not accept arguments" in result.stderr
-7
View File
@@ -1,7 +0,0 @@
"""
foobuzz test
"""
def test_foo():
pass
+1 -1
View File
@@ -15,7 +15,7 @@ from nexus_sync.common import (
) )
def test_heartbeat_request_accepts_current_mvp_payload() -> None: def test_heartbeat_request_accepts_payload() -> None:
payload = HeartbeatRequest( payload = HeartbeatRequest(
client_id="macbook-pro-01", client_id="macbook-pro-01",
observed_at=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC), observed_at=datetime(2026, 5, 24, 13, 20, 30, tzinfo=UTC),
+174
View File
@@ -0,0 +1,174 @@
from datetime import UTC, datetime
from fastapi.testclient import TestClient
from nexus_sync.common import Command, CommandKind, CommandResultStatus, CommandStatus
from nexus_sync.server import (
DEFAULT_COMMAND_POLL_SECONDS,
DEFAULT_IDLE_POLL_SECONDS,
Store,
InMemoryStore,
create_app,
)
def _heartbeat_payload(client_id: str = "macbook-pro-01", result: dict | None = None) -> dict:
return {
"client_id": client_id,
"observed_at": "2026-05-24T13:20:30Z",
"client": {
"hostname": "macbook-pro.local",
"platform": "darwin",
"version": "0.1.0",
},
"state": {
"local_time": "2026-05-24T16:20:30+03:00",
"uptime_seconds": 1200,
},
"last_command_result": result,
}
def _client(store: Store | None = None) -> TestClient:
app = create_app(
store=store or InMemoryStore(),
client_tokens={"macbook-pro-01": "client-token", "other-client": "other-token"},
)
return TestClient(app)
def test_heartbeat_requires_bearer_token() -> None:
client = _client()
response = client.post("/api/v1/client/heartbeat", json=_heartbeat_payload())
assert response.status_code == 401
def test_heartbeat_rejects_token_for_another_client() -> None:
client = _client()
response = client.post(
"/api/v1/client/heartbeat",
json=_heartbeat_payload(),
headers={"Authorization": "Bearer other-token"},
)
assert response.status_code == 403
def test_heartbeat_returns_bad_request_for_invalid_payload() -> None:
payload = _heartbeat_payload()
payload["unexpected"] = True
client = _client()
response = client.post(
"/api/v1/client/heartbeat",
json=payload,
headers={"Authorization": "Bearer client-token"},
)
assert response.status_code == 400
def test_heartbeat_accepts_state_without_command() -> None:
store = InMemoryStore()
client = _client(store)
response = client.post(
"/api/v1/client/heartbeat",
json=_heartbeat_payload(),
headers={"Authorization": "Bearer client-token"},
)
assert response.status_code == 200
assert response.json()["command"] is None
assert response.json()["next_poll_after_seconds"] == DEFAULT_IDLE_POLL_SECONDS
assert store.clients["macbook-pro-01"].hostname == "macbook-pro.local"
def test_heartbeat_delivers_pending_command_and_marks_it_delivered() -> None:
store = InMemoryStore()
store.enqueue_command(
Command(
id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B",
kind=CommandKind.EXEC,
name="hostname",
args={},
timeout_seconds=30,
),
client_id="macbook-pro-01",
now=datetime(2026, 5, 24, 13, 20, 0, tzinfo=UTC),
)
client = _client(store)
response = client.post(
"/api/v1/client/heartbeat",
json=_heartbeat_payload(),
headers={"Authorization": "Bearer client-token"},
)
body = response.json()
assert response.status_code == 200
assert body["next_poll_after_seconds"] == DEFAULT_COMMAND_POLL_SECONDS
assert body["command"]["id"] == "cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"
assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].status == CommandStatus.DELIVERED
assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].attempts == 1
def test_heartbeat_records_command_result() -> None:
store = InMemoryStore()
store.enqueue_command(
Command(
id="cmd_01JY3H8V8W8P3FXDR3S2BM7M6B",
kind=CommandKind.EXEC,
name="hostname",
args={},
timeout_seconds=30,
),
client_id="macbook-pro-01",
now=datetime(2026, 5, 24, 13, 20, 0, tzinfo=UTC),
)
store.take_next_command("macbook-pro-01", datetime(2026, 5, 24, 13, 20, 1, tzinfo=UTC))
client = _client(store)
response = client.post(
"/api/v1/client/heartbeat",
json=_heartbeat_payload(
result={
"command_id": "cmd_01JY3H8V8W8P3FXDR3S2BM7M6B",
"status": "succeeded",
"started_at": "2026-05-24T13:20:02Z",
"finished_at": "2026-05-24T13:20:03Z",
"return_code": 0,
"stdout": "host\n",
"stderr": "",
}
),
headers={"Authorization": "Bearer client-token"},
)
assert response.status_code == 200
assert store.commands["cmd_01JY3H8V8W8P3FXDR3S2BM7M6B"].status == CommandStatus.SUCCEEDED
assert len(store.results) == 1
assert store.results[0].status == CommandResultStatus.SUCCEEDED
def test_heartbeat_rejects_unknown_command_result() -> None:
client = _client()
response = client.post(
"/api/v1/client/heartbeat",
json=_heartbeat_payload(
result={
"command_id": "cmd_unknown",
"status": "succeeded",
"return_code": 0,
"stdout": "",
"stderr": "",
}
),
headers={"Authorization": "Bearer client-token"},
)
assert response.status_code == 409