Metadata-Version: 2.4
Name: puda
Version: 0.0.19rc1
Summary: Official Python SDK for PUDA, the runtime environment for Physical AI
Author: zhao
Author-email: zhao <20024592+agentzhao@users.noreply.github.com>
License-Expression: MIT
Requires-Dist: cryptography>=44.0.0
Requires-Dist: nats-py>=2.12.0
Requires-Dist: psutil>=6.1.1
Requires-Dist: pydantic>=2.12.5
Requires-Python: >=3.10
Description-Content-Type: text/markdown

# Puda Comms

A Python module for communication between machines and command services via NATS messaging. Provides client-side services for sending commands, machine-side clients for receiving commands, and data models for structured message exchange.

## Overview

The `puda` module enables asynchronous, reliable communication between command services and machines using NATS (NATS JetStream for guaranteed delivery). It handles:

- **Command execution**: Send commands to machines and receive responses
- **Message routing**: Queue commands (sequential execution) and immediate commands (control operations)
- **State management**: Thread-safe execution state tracking for cancellation and locking
- **Connection management**: Automatic NATS connection handling with async context managers

## Components

The module consists of four main components:

### 1. Models (`models.py`)

Data models for structured message exchange. All models use Pydantic for validation and serialization.

#### Enums

##### `CommandResponseStatus`
Status of a command response:
- `SUCCESS`: Command executed successfully
- `ERROR`: Command execution failed

##### `CommandResponseCode`
Error codes for command responses:
- `COMMAND_CANCELLED`: Command was cancelled before completion
- `JSON_DECODE_ERROR`: Failed to decode JSON payload
- `EXECUTION_ERROR`: General execution error
- `EXECUTION_LOCKED`: Execution is locked (another command is running)
- `UNKNOWN_COMMAND`: Command name not recognized
- `PAUSE_ERROR`: Error occurred while pausing execution
- `RESUME_ERROR`: Error occurred while resuming execution
- `NO_EXECUTION`: No execution found
- `RUN_ID_MISMATCH`: Run ID doesn't match current execution
- `CANCEL_ERROR`: Error occurred while cancelling execution
- `MACHINE_PAUSED`: Machine is currently paused

##### `MessageType`
Type of NATS message:
- `COMMAND`: Command message sent to machine
- `RESPONSE`: Response message from machine
- `LOG`: Log message
- `ALERT`: Alert message
- `MEDIA`: Media message

##### `ImmediateCommand`
Command names for immediate/control commands:
- `PAUSE`: Pause the current execution
- `RESUME`: Resume a paused execution
- `CANCEL`: Cancel the current execution

#### Data Models

##### `CommandRequest`
Represents a command to be sent to a machine.

**Fields:**
- `name` (str): The command name to execute
- `machine_id` (str): Machine ID to send the command to (required)
- `params` (Dict[str, Any]): Command parameters (default: empty dict)
- `step_number` (int): Execution step number for tracking progress
- `version` (str): Command version (default: "1.0")

**Example:**
```python
command = CommandRequest(
    name="attach_tip",
    machine_id="first",
    params={"deck_slot": "A3", "well_name": "G8"},
    step_number=2,
    version="1.0"
)
```

##### `CommandResponse`
Represents the result of a command execution.

**Fields:**
- `status` (CommandResponseStatus): Status of the command response (SUCCESS or ERROR)
- `completed_at` (str): ISO 8601 UTC timestamp (auto-generated)
- `code` (Optional[str]): Error code if status is ERROR
- `message` (Optional[str]): Human-readable error message

**Example:**
```python
response = CommandResponse(
    status=CommandResponseStatus.SUCCESS,
    completed_at="2026-01-20T02:00:46Z"
)
```

**Error Example:**
```python
error_response = CommandResponse(
    status=CommandResponseStatus.ERROR,
    code="EXECUTION_ERROR",
    message="Failed to attach tip: deck_slot A3 not found",
    completed_at="2026-01-20T02:00:46Z"
)
```

##### `MessageHeader`
Header metadata for NATS messages.

**Fields:**
- `message_type` (MessageType): Type of message (COMMAND, RESPONSE, LOG, etc.)
- `version` (str): Message version (default: "1.0")
- `timestamp` (str): ISO 8601 UTC timestamp (auto-generated)
- `user_id` (str): User ID who initiated the command
- `username` (str): Username who initiated the command
- `machine_id` (str): Identifier for the target machine
- `run_id` (Optional[str]): Unique identifier (UUID) for the run/workflow

**Example:**
```python
header = MessageHeader(
    message_type=MessageType.RESPONSE,
    version="1.0",
    timestamp="2026-01-20T02:00:46Z",
    user_id="user123",
    username="John Doe",
    machine_id="first",
    run_id="092073e6-13d0-4756-8d99-eff1612a5a72"
)
```

##### `NATSMessage`
Complete NATS message structure combining header with optional command or response data.

**Fields:**
- `header` (MessageHeader): Message header (required)
- `command` (Optional[CommandRequest]): Command request (for command messages)
- `response` (Optional[CommandResponse]): Command response (for response messages)

**Structure:**
- For command messages: include `header` with `message_type=COMMAND` and `command` field
- For response messages: include `header` with `message_type=RESPONSE` and `response` field

**Complete Message Example:**
```json
{
  "header": {
    "message_type": "response",
    "version": "1.0",
    "timestamp": "2026-01-20T02:00:46Z",
    "user_id": "user123",
    "username": "John Doe",
    "machine_id": "first",
    "run_id": "092073e6-13d0-4756-8d99-eff1612a5a72"
  },
  "command": {
    "name": "attach_tip",
    "params": {
      "deck_slot": "A3",
      "well_name": "G8"
    },
    "step_number": 2,
    "version": "1.0"
  },
  "response": {
    "status": "success",
    "completed_at": "2026-01-20T02:00:46Z",
    "code": null,
    "message": null
  }
}
```

### 2. CommandService (`command_service.py`)

Client-side service for sending commands to machines via NATS. Handles:
- Connecting to NATS servers
- Sending commands to machines (queue or immediate)
- Waiting for and handling responses
- Managing command lifecycle (run_id, step_number, etc.)
- Automatic connection cleanup via async context manager

See [Sending Commands](#sending-commands) section for usage examples.

### 3. EdgeNatsClient (`machine_client.py`)

Basic default NATS client for generic machines. Handles commands, telemetry, and events following the `puda.{machine_id}.{category}.{sub_category}` pattern. Provides:
- Subscribing to command streams (queue and immediate) via JetStream with exactly-once delivery
- Processing incoming commands and sending command responses
- Publishing telemetry (core NATS, no JetStream)
- Publishing events (core NATS, fire-and-forget)
- Responding to direct and fleet-wide Core NATS ping requests
- Connection management and reconnection handling

**Note:** This is a generic client. Machine-specific methods should be implemented in the machine-edge client.

#### Core NATS ping/pong

`EdgeRunner` automatically subscribes each connected edge to both:

```text
puda.<machine_id>.cmd.ping
puda.cmd.ping
```

A Core NATS request with payload `ping` receives structured JSON:

```json
{
  "status": "pong",
  "machine_id": "first",
  "timestamp": "2026-08-27T07:23:56Z",
  "sdk_version": "0.0.19-rc1",
  "uptime_seconds": 12.5,
  "run_status": "idle",
  "description": "Cartesian gantry for well-plate liquid handling.",
  "tlm_streams": [
    {
      "name": "pos",
      "subject": "puda.first.tlm.stream.pos",
      "interval": 3.0,
      "description": "Current cartesian position."
    }
  ],
  "local_ip": "192.168.1.10",
  "tailscale_ip": "100.99.243.61",
  "magicdns": "host.tailnet.ts.net"
}
```

`run_status` is derived from the SDK's in-memory execution lock: `busy` while
a command is executing and `idle` otherwise. It is intentionally not read from
or persisted to machine KV state.

`description` is the first paragraph of the driver class docstring, collapsed
to a single line. `EdgeRunner` copies it onto the NATS client at startup. Pass
`EdgeNatsClient(..., description="...")` to override it. The field is omitted
when unset.

`tlm_streams` lists every telemetry stream the edge has declared, with its
subject, interval in seconds, and description. It is always present and may be
empty. See [Telemetry streams](#telemetry-streams).

`local_ip`, `tailscale_ip`, and `magicdns` are best-effort host addresses.
`local_ip` is the LAN IPv4 (default-route address, falling back if Tailscale
is the default route). `tailscale_ip` and `magicdns` come from
`tailscale status --json`. Each field is omitted when it cannot be discovered.

Ping is intentionally Core NATS request/reply, not a durable JetStream
immediate command. It reports whether the edge is responsive now and does not
create an offline backlog.

`puda machine list` sends one broadcast request to `puda.cmd.ping`, gathers all
pong replies during its discovery window, deduplicates them by `machine_id`,
and lists only edges that are responsive at that moment. Each reply's
`description` is included so agents can tell what a machine does without
fetching the command catalog. Host address fields (`local_ip`,
`tailscale_ip`, `magicdns`) are forwarded the same way so a caller can reach
the host over LAN or Tailscale. The CLI also reports `livestream_count`: how
many fleet `LIVESTREAMS` registry records attach to that machine. That count
is registered PUDA livestreams only; unregistered cameras on the host may
exist and are not included. Use `puda livestream list --machines <id>` for
registered names, hosts, and URLs. It also reports `tlm_stream_count`.
`puda machine info` (formerly `puda machine ping`, which still works as an
alias) shows the full `tlm_streams` list and joins the full
registered records (host, stream name, derived protocol URLs). Those records
store host and stream name; protocol URLs are derived and are not part of the
edge pong payload.

### 4. EdgeRunner (`edge_runner.py`)

Dispatches incoming NATS commands onto the machine driver. Only methods marked with `@command` are advertised and executed.

Optional `@safety` metadata is published with the command catalog so agents can prompt the operator before execution. It does not change dispatch.

```python
from puda import command, safety

@command
@safety(
    summary="Collision risk. Confirm the deck is clear and the machine is homed.",
    confirm=True,
)
def move(self, x: float, y: float, z: float) -> dict:
    """Move to an absolute position."""
    ...
```

`summary`, `hazards`, `requires`, `forbidden_when`, and `confirm` are the only `@safety` fields, and they are keyword-only. `confirm` defaults to `false`. Agents must prompt the user before executing a tagged command only when the driver sets `confirm=True`. `hazards`, `requires`, and `forbidden_when` are optional natural-language context.

**Failure semantics:** Command handlers must raise an exception to indicate failure. Returning `False` (or any other value) is still a successful PUDA response. A `False` return is serialized as:

```json
{"result": false}
```

#### Telemetry streams

Mark a driver method with `@tlm_stream(interval=...)` to publish its return value as a telemetry stream. `EdgeRunner` calls the method every `interval` seconds, in a worker thread as it does for commands, and publishes the result to `puda.<machine_id>.tlm.stream.<name>`. Each stream runs at its own rate and can run at the same time as a command.

```python
from puda import command, tlm_stream

@tlm_stream(interval=0.25)
def weight(self) -> dict | None:
    """Mass in grams."""
    return {"g": self._latest, "fresh": self._fresh}

@command
@tlm_stream(interval=3.0, name="pos")
def get_position(self) -> dict[str, float]:
    ...
```

- The stream name defaults to the method name. Names may contain letters, digits, `_` and `-`.
- The first paragraph of the method docstring is advertised as the stream description.
- Return `None` to skip a sample. Exceptions are logged and the stream keeps running.
- Subscribe to `puda.<machine_id>.tlm.stream.>` for every stream on a machine.

For data that arrives as events from a driver thread rather than on a fixed schedule, call `nats_client.publish_tlm_stream_threadsafe("<name>", data)`.

#### Machine state, heartbeat, and host health

Mark one driver method with `@machine_state` to add fields to every `MACHINE_STATE` update:

```python
from puda import machine_state

@machine_state
def snapshot(self) -> dict:
    return {"homed": self._homed}
```

The method is called whenever machine state is published: at startup and reconnect, when a queued command starts and ends, after immediate commands, on errors, and at shutdown. It is not polled, so an idle machine's state is written once. It runs on the event loop rather than in a worker thread, so return cached values; a blocking device read stalls command handling and telemetry. Use `@tlm_stream` for values that change outside commands. Don't return `state`, `run_id`, or `timestamp`, which PUDA sets.

`EdgeRunner` also publishes the heartbeat (every 5 seconds) and host health (`cpu`, `mem` and `temp` on `tlm.health`, every 5 seconds) without any edge code. A minimal edge is:

```python
runner = EdgeRunner(nats_client=nats_client, machine_driver=driver)
```

### 5. ExecutionState (`execution_state.py`)

Thread-safe state management for command execution. Provides:
- Execution lock to prevent concurrent commands
- Current task tracking for cancellation
- Run ID matching for cancel operations
- Thread-safe access to execution state

## Sending Commands

The `CommandService` provides a high-level interface for sending commands to machines via NATS. See [`tests/commands.py`](tests/commands.py) and [`tests/batch_commands.py`](tests/batch_commands.py) for complete examples.

### Recommended Usage: Async Context Manager

The recommended way to use `CommandService` is with an async context manager, which automatically handles connection and disconnection. See [`tests/commands.py`](tests/commands.py) for complete examples.

### Command Types

#### Queue Commands

Queue commands are regular commands that are executed in sequence. Use `send_queue_command()` for machine-specific operations.

**Note:** Available commands depend on the machine you are controlling. Different machines support different command sets (e.g., `first` machine supports commands like `load_deck`, `attach_tip`, `aspirate_from`, `dispense_to`, `drop_tip`, etc.).

Both `send_queue_command()`, `send_queue_commands()`, and `send_immediate_command()` accept an optional `timeout` parameter (default: 120 seconds):

```python
# Single command (machine_id must be in CommandRequest)
reply = await service.send_queue_command(
    request=request,  # request.machine_id must be set
    run_id=run_id,
    user_id="user123",
    username="John Doe",
    timeout=60  # Wait up to 60 seconds
)

# Multiple commands (timeout applies to each command)
# Each command in the list must have machine_id set
reply = await service.send_queue_commands(
    requests=commands,  # Each CommandRequest must have machine_id
    run_id=run_id,
    user_id="user123",
    username="John Doe",
    timeout=60  # Wait up to 60 seconds per command
)
```

**Examples:**

See [`tests/commands.py`](tests/commands.py) for complete examples.

#### Immediate Commands

Immediate commands are control commands that interrupt or modify execution. Use `send_immediate_command()` for:
- `pause`: Pause the current execution
- `resume`: Resume a paused execution
- `cancel`: Cancel the current execution

**Examples:**

See [`tests/commands.py`](tests/commands.py) for complete examples.



### Sending Command Sequences

You can send multiple commands in sequence using `send_queue_commands()`, which sends commands one by one and waits for each response before sending the next. If any command fails or times out, it stops immediately and returns the error response.

**Loading Commands from JSON (Recommended for LLM-generated commands):**

When generating commands from an LLM or loading from external sources, you can store commands in a JSON file and load them. See [`tests/batch_commands.py`](tests/batch_commands.py) for a complete example.

### Error Handling

Always check the response status and handle errors appropriately:

```python
reply: NATSMessage = await service.send_queue_command(
    request=request,  # request.machine_id must be set
    run_id=run_id,
    user_id="user123",
    username="John Doe"
)

if reply is None:
    # Command timed out or failed to send
    logger.error("Command failed or timed out")
elif reply.response is not None and reply.response.status == CommandResponseStatus.SUCCESS:
    # Command succeeded
    logger.info("Command completed successfully")
else:
    # Command failed with error
    logger.error("Command failed with code: %s, message: %s", 
                reply.response.code if reply.response else None,
                reply.response.message if reply.response else None)
```

### Configuration

#### NATS Server Configuration

The `CommandService` requires NATS server URLs to be specified explicitly. There are no default values. You must provide servers in one of two ways:

**Option 1: Via environment variable (comma-separated string)**

Set the `NATS_SERVERS` environment variable with comma-separated server URLs:

```bash
export NATS_SERVERS="nats://192.168.50.201:4222,nats://192.168.50.201:4223,nats://192.168.50.201:4224"
```

Then parse it when creating a `CommandService`:
```python
import os
nats_servers = [s.strip() for s in os.getenv("NATS_SERVERS", "").split(",") if s.strip()]
service = CommandService(servers=nats_servers)
```

**Option 2: Directly as a list**

Specify servers directly when creating a `CommandService`:
```python
service = CommandService(servers=["nats://192.168.50.201:4222", "nats://192.168.50.201:4223", "nats://192.168.50.201:4224"])
```
## Validation

All models use Pydantic for validation, ensuring:
- Type checking for all fields
- Required fields are present
- Default values are applied correctly
- JSON serialization/deserialization works correctly
