puda-comms 0.0.3__tar.gz → 0.0.5__tar.gz

{puda_comms-0.0.3 → puda_comms-0.0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: puda-comms
-Version: 0.0.3
+Version: 0.0.5
 Summary: Communication library for the PUDA platform.
 Author: zhao
 Author-email: zhao <20024592+agentzhao@users.noreply.github.com>
@@ -121,6 +121,8 @@ Header metadata for NATS messages.
 - `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
 
@@ -130,6 +132,8 @@ 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"
 )
@@ -154,6 +158,8 @@ Complete NATS message structure combining header with optional command or respon
     "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"
   },
@@ -229,6 +235,8 @@ reply = await service.send_queue_command(
     request=request,
     machine_id="first",
     run_id=run_id,
+    user_id="user123",
+    username="John Doe",
     timeout=60  # Wait up to 60 seconds
 )
 
@@ -237,6 +245,8 @@ reply = await service.send_queue_commands(
     requests=commands,
     machine_id="first",
     run_id=run_id,
+    user_id="user123",
+    username="John Doe",
     timeout=60  # Wait up to 60 seconds per command
 )
 ```
@@ -274,7 +284,9 @@ Always check the response status and handle errors appropriately:
 reply: NATSMessage = await service.send_queue_command(
     request=request,
     machine_id="first",
-    run_id=run_id
+    run_id=run_id,
+    user_id="user123",
+    username="John Doe"
 )
 
 if reply is None:

{puda_comms-0.0.3 → puda_comms-0.0.5}/README.md

@@ -109,6 +109,8 @@ Header metadata for NATS messages.
 - `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
 
@@ -118,6 +120,8 @@ 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"
 )
@@ -142,6 +146,8 @@ Complete NATS message structure combining header with optional command or respon
     "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"
   },
@@ -217,6 +223,8 @@ reply = await service.send_queue_command(
     request=request,
     machine_id="first",
     run_id=run_id,
+    user_id="user123",
+    username="John Doe",
     timeout=60  # Wait up to 60 seconds
 )
 
@@ -225,6 +233,8 @@ reply = await service.send_queue_commands(
     requests=commands,
     machine_id="first",
     run_id=run_id,
+    user_id="user123",
+    username="John Doe",
     timeout=60  # Wait up to 60 seconds per command
 )
 ```
@@ -262,7 +272,9 @@ Always check the response status and handle errors appropriately:
 reply: NATSMessage = await service.send_queue_command(
     request=request,
     machine_id="first",
-    run_id=run_id
+    run_id=run_id,
+    user_id="user123",
+    username="John Doe"
 )
 
 if reply is None:

{puda_comms-0.0.3 → puda_comms-0.0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "puda-comms"
-version = "0.0.3"
+version = "0.0.5"
 description = "Communication library for the PUDA platform."
 readme = "README.md"
 authors = [
@@ -17,6 +17,17 @@ dependencies = [
 [tool.uv.sources]
 puda-drivers = {workspace = true}
 
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint.mccabe]
+max-complexity = 10
+
 [build-system]
 requires = ["uv_build>=0.9.18,<0.10.0"]
 build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "ruff>=0.14.13",
+]

{puda_comms-0.0.3 → puda_comms-0.0.5}/src/puda_comms/command_service.py

@@ -12,11 +12,17 @@ import json
 import logging
 import signal
 from datetime import datetime, timezone
-from typing import Dict, Any, Optional, Tuple
+from typing import Dict, Any, Optional
 import nats
 from nats.js.client import JetStreamContext
 from nats.aio.msg import Msg
-from puda_comms.models import CommandRequest, CommandResponse, CommandResponseStatus, NATSMessage, MessageHeader, MessageType
+from puda_comms.models import (
+    CommandRequest,
+    CommandResponseStatus,
+    NATSMessage,
+    MessageHeader,
+    MessageType,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -37,7 +43,7 @@ class ResponseHandler:
     def __init__(self, js: JetStreamContext, machine_id: str):
         self.js = js
         self.machine_id = machine_id
-        self._pending_responses: Dict[str, Tuple[asyncio.Event, CommandResponse]] = {}
+        self._pending_responses: Dict[str, Dict[str, Any]] = {}  # {'event': asyncio.Event, 'response': Optional[NATSMessage]}
         self._queue_consumer = None
         self._immediate_consumer = None
         self._initialized = False
@@ -102,8 +108,8 @@ class ResponseHandler:
                 
                 # Get the pending response
                 pending = self._pending_responses[key]
-                # Store the full NATSMessage JSON structure
-                pending['response'] = message.model_dump()
+                # Store the NATSMessage directly
+                pending['response'] = message
                 # Signal that response was received
                 # Don't delete here - let get_response() delete it after retrieval
                 pending['event'].set()
@@ -152,7 +158,7 @@ class ResponseHandler:
         }
         return event
     
-    def get_response(self, run_id: str, step_number: int) -> Optional[Dict[str, Any]]:
+    def get_response(self, run_id: str, step_number: int) -> Optional[NATSMessage]:
         """
         Get the response for a pending command.
         
@@ -161,7 +167,7 @@ class ResponseHandler:
             step_number: Step number for the command
         
         Returns:
-            The NATSMessage dict structure if available, None otherwise
+            The NATSMessage if available, None otherwise
         """
         key = f"{run_id}:{str(step_number)}"
         if key in self._pending_responses:
@@ -266,9 +272,9 @@ class CommandService:
         max_attempts = 3
         connect_timeout = 3  # 3 seconds timeout per connection attempt
         
-        for attempt in range(1, max_attempts + 1):
+        for attempt in range(max_attempts):
             try:
-                logger.info("Connection attempt %d/%d to NATS servers: %s", attempt, max_attempts, self.servers)
+                logger.info("Connection attempt %d/%d to NATS servers: %s", attempt + 1, max_attempts, self.servers)
                 self.nc = await asyncio.wait_for(
                     nats.connect(
                         servers=self.servers,
@@ -285,14 +291,14 @@ class CommandService:
                 return True
                 
             except asyncio.TimeoutError:
-                logger.warning("Connection attempt %d/%d timed out after %d seconds", attempt, max_attempts, connect_timeout)
-                if attempt < max_attempts:
+                logger.warning("Connection attempt %d/%d timed out after %d seconds", attempt + 1, max_attempts, connect_timeout)
+                if attempt < max_attempts - 1:
                     logger.info("Retrying connection...")
                 else:
                     logger.error("Failed to connect after %d attempts. Giving up.", max_attempts)
             except Exception as e:
-                logger.warning("Connection attempt %d/%d failed: %s", attempt, max_attempts, e)
-                if attempt < max_attempts:
+                logger.warning("Connection attempt %d/%d failed: %s", attempt + 1, max_attempts, e)
+                if attempt < max_attempts - 1:
                     logger.info("Retrying connection...")
                 else:
                     logger.error("Failed to connect after %d attempts. Giving up.", max_attempts)
@@ -343,6 +349,8 @@ class CommandService:
         request: CommandRequest,
         machine_id: str,
         run_id: str,
+        user_id: str,
+        username: str,
         timeout: int = 120
     ) -> Optional[NATSMessage]:
         """
@@ -352,6 +360,8 @@ class CommandService:
             request: CommandRequest model containing command details
             machine_id: Machine ID to send the command to
             run_id: Run ID for the command
+            user_id: User ID who initiated the command
+            username: Username who initiated the command
             timeout: Maximum time to wait for response in seconds
         
         Returns:
@@ -364,8 +374,8 @@ class CommandService:
         subject = f"{NAMESPACE}.{machine_id}.cmd.queue"
         
         logger.info(
-            "Sending queue command: machine_id=%s, command=%s, run_id=%s, step_number=%s",
-            machine_id, request.name, run_id, request.step_number
+            "Sending queue command: subject=%s, command=%s, run_id=%s, step_number=%s",
+            subject, request.name, run_id, request.step_number
         )
         
         # Get or create response handler for this machine
@@ -374,7 +384,7 @@ class CommandService:
         response_event = response_handler.register_pending(run_id, request.step_number)
         
         # Build payload
-        payload = self._build_command_payload(request, machine_id, run_id)
+        payload = self._build_command_payload(request, machine_id, run_id, user_id, username)
 
         try:
             # Publish to JetStream
@@ -397,36 +407,107 @@ class CommandService:
             await asyncio.sleep(0.1)
             
             # Get the response
-            response_data = response_handler.get_response(run_id, request.step_number)
-            if response_data is None:
-                return None
-            
-            return NATSMessage.model_validate(response_data)
+            return response_handler.get_response(run_id, request.step_number)
 
         except Exception as e:
             logger.error("Error sending queue command: %s", e)
             response_handler.remove_pending(run_id, request.step_number)
             return None
     
+    async def start_run(
+        self,
+        machine_id: str,
+        run_id: str,
+        user_id: str,
+        username: str,
+        timeout: int = 120
+    ) -> Optional[NATSMessage]:
+        """
+        Send START immediate command to begin a run.
+        
+        Args:
+            machine_id: Machine ID to send the command to
+            run_id: Run ID for the command
+            user_id: User ID who initiated the command
+            username: Username who initiated the command
+            timeout: Maximum time to wait for response in seconds
+        
+        Returns:
+            NATSMessage if successful, None if failed or timeout
+        """
+        request = CommandRequest(
+            name="start",
+            params={},
+            step_number=0
+        )
+        return await self.send_immediate_command(
+            request=request,
+            machine_id=machine_id,
+            run_id=run_id,
+            user_id=user_id,
+            username=username,
+            timeout=timeout
+        )
+    
+    async def complete_run(
+        self,
+        machine_id: str,
+        run_id: str,
+        user_id: str,
+        username: str,
+        timeout: int = 120
+    ) -> Optional[NATSMessage]:
+        """
+        Send COMPLETE immediate command to end a run.
+        
+        Args:
+            machine_id: Machine ID to send the command to
+            run_id: Run ID for the command
+            user_id: User ID who initiated the command
+            username: Username who initiated the command
+            timeout: Maximum time to wait for response in seconds
+        
+        Returns:
+            NATSMessage if successful, None if failed or timeout
+        """
+        request = CommandRequest(
+            name="complete",
+            params={},
+            step_number=0
+        )
+        return await self.send_immediate_command(
+            request=request,
+            machine_id=machine_id,
+            run_id=run_id,
+            user_id=user_id,
+            username=username,
+            timeout=timeout
+        )
+    
     async def send_queue_commands(
         self,
         *,
         requests: list[CommandRequest],
         machine_id: str,
         run_id: str,
+        user_id: str,
+        username: str,
         timeout: int = 120
     ) -> Optional[NATSMessage]:
         """
         Send multiple queue commands sequentially and wait for responses.
         
-        Sends commands one by one, waiting for each response before sending the next.
-        If any command fails or times out, stops immediately and returns the error response.
-        If all commands succeed, returns the last command's response.
+        Automatically sends START command before the sequence and COMPLETE command after
+        successful completion. Sends commands one by one, waiting for each response before
+        sending the next. If any command fails or times out, stops immediately and returns
+        the error response. If all commands succeed, returns the last command's response.
         
         Args:
             requests: List of CommandRequest models to send sequentially
             machine_id: Machine ID to send the commands to
             run_id: Run ID for all commands
+            user_id: User ID who initiated the commands
+            username: Username who initiated the commands
             timeout: Maximum time to wait for each response in seconds
         
         Returns:
@@ -447,74 +528,131 @@ class CommandService:
             run_id
         )
         
+        # Always send START command before sequence
+        logger.info("Sending START command before sequence")
+        start_response = await self.start_run(
+            machine_id=machine_id,
+            run_id=run_id,
+            user_id=user_id,
+            username=username,
+            timeout=timeout
+        )
+        if start_response is None:
+            logger.error("START command timed out")
+            return None
+        if start_response.response and start_response.response.status == CommandResponseStatus.ERROR:
+            logger.error("START command failed: %s", start_response.response.message)
+            return start_response
+        
         last_response: Optional[NATSMessage] = None
         
-        for idx, request in enumerate(requests, start=1):
-            logger.info(
-                "Sending command %d/%d: %s (step %s)",
-                idx,
-                len(requests),
-                request.name,
-                request.step_number
-            )
-            
-            response = await self.send_queue_command(
-                request=request,
-                machine_id=machine_id,
-                run_id=run_id,
-                timeout=timeout
-            )
-            
-            # Check if command failed (None means timeout or exception)
-            if response is None:
-                logger.error(
-                    "Command %d/%d failed or timed out: %s (step %s)",
+        try:
+            for idx, request in enumerate(requests, start=1):
+                # Validate request - convert dict to CommandRequest if needed
+                if isinstance(request, dict):
+                    request = CommandRequest.model_validate(request)
+                elif not isinstance(request, CommandRequest):
+                    raise ValueError(f"Request {idx} must be a CommandRequest or dict, got {type(request)}")
+                
+                logger.info(
+                    "Sending command %d/%d: %s (step %s)",
                     idx,
                     len(requests),
                     request.name,
                     request.step_number
                 )
-                return None
             
-            # Check if command returned an error status
-            if response.response is not None:
-                if response.response.status == CommandResponseStatus.ERROR:
+                response = await self.send_queue_command(
+                    request=request,
+                    machine_id=machine_id,
+                    run_id=run_id,
+                    user_id=user_id,
+                    username=username,
+                    timeout=timeout
+                )
+            
+                # Check if command failed (None means timeout or exception)
+                if response is None:
                     logger.error(
-                        "Command %d/%d failed with error: %s (step %s) - code: %s, message: %s",
+                        "Command %d/%d failed or timed out: %s (step %s)",
+                        idx,
+                        len(requests),
+                        request.name,
+                        request.step_number
+                    )
+                    return None
+            
+                # Check if command returned an error status
+                if response.response is not None:
+                    if response.response.status == CommandResponseStatus.ERROR:
+                        logger.error(
+                            "Command %d/%d failed with error: %s (step %s) - code: %s, message: %s",
+                            idx,
+                            len(requests),
+                            request.name,
+                            request.step_number,
+                            response.response.code,
+                            response.response.message
+                        )
+                        return response
+                    
+                    # Command succeeded, store as last response
+                    last_response = response
+                    logger.info(
+                        "Command %d/%d succeeded: %s (step %s)",
                         idx,
                         len(requests),
                         request.name,
-                        request.step_number,
-                        response.response.code,
-                        response.response.message
+                        request.step_number
+                    )
+                else:
+                    # Response exists but has no response data (shouldn't happen, but handle it)
+                    logger.warning(
+                        "Command %d/%d returned response with no response data: %s (step %s)",
+                        idx,
+                        len(requests),
+                        request.name,
+                        request.step_number
                     )
                     return response
-                
-                # Command succeeded, store as last response
-                last_response = response
-                logger.info(
-                    "Command %d/%d succeeded: %s (step %s)",
-                    idx,
-                    len(requests),
-                    request.name,
-                    request.step_number
-                )
-            else:
-                # Response exists but has no response data (shouldn't happen, but handle it)
-                logger.warning(
-                    "Command %d/%d returned response with no response data: %s (step %s)",
-                    idx,
-                    len(requests),
-                    request.name,
-                    request.step_number
+            
+            logger.info(
+                "All %d commands completed successfully",
+                len(requests)
+            )
+            
+            # Always send COMPLETE command after successful sequence
+            logger.info("Sending COMPLETE command after successful sequence")
+            complete_response = await self.complete_run(
+                machine_id=machine_id,
+                run_id=run_id,
+                user_id=user_id,
+                username=username,
+                timeout=timeout
+            )
+            if complete_response is None:
+                logger.error("COMPLETE command timed out")
+                return None
+            if complete_response.response and complete_response.response.status == CommandResponseStatus.ERROR:
+                logger.error("COMPLETE command failed: %s", complete_response.response.message)
+                return complete_response
+            # Return the last command response, not the COMPLETE response
+            return last_response
+        except Exception as e:
+            # If any error occurs during command execution, try to complete the run
+            # to clean up state (but don't fail if this also fails)
+            logger.warning("Error during command sequence, attempting to complete run: %s", e)
+            try:
+                await self.complete_run(
+                    machine_id=machine_id,
+                    run_id=run_id,
+                    user_id=user_id,
+                    username=username,
+                    timeout=timeout
                 )
-                return response
-        
-        logger.info(
-            "All %d commands completed successfully",
-            len(requests)
-        )
-        return last_response
+            except Exception as cleanup_error:
+                logger.error("Failed to complete run during error cleanup: %s", cleanup_error)
+            raise
     
     async def send_immediate_command(
         self,
@@ -522,6 +660,8 @@ class CommandService:
         request: CommandRequest,
         machine_id: str,
         run_id: str,
+        user_id: str,
+        username: str,
         timeout: int = 120
     ) -> Optional[NATSMessage]:
         """
@@ -531,6 +671,8 @@ class CommandService:
             request: CommandRequest model containing command details
             machine_id: Machine ID to send the command to
             run_id: Run ID for the command
+            user_id: User ID who initiated the command
+            username: Username who initiated the command
             timeout: Maximum time to wait for response in seconds
         
         Returns:
@@ -555,7 +697,7 @@ class CommandService:
         response_received = response_handler.register_pending(run_id, request.step_number)
         
         # Build payload
-        payload = self._build_command_payload(request, machine_id, run_id)
+        payload = self._build_command_payload(request, machine_id, run_id, user_id, username)
 
         try:
             # Publish to JetStream
@@ -578,11 +720,7 @@ class CommandService:
             await asyncio.sleep(0.1)
             
             # Get the response
-            response_data = response_handler.get_response(run_id, request.step_number)
-            if response_data is None:
-                return None
-            
-            return NATSMessage.model_validate(response_data)
+            return response_handler.get_response(run_id, request.step_number)
             
         except Exception as e:
             logger.error("Error sending immediate command: %s", e)
@@ -635,7 +773,9 @@ class CommandService:
         self,
         command_request: CommandRequest,
         machine_id: str,
-        run_id: str
+        run_id: str,
+        user_id: str,
+        username: str
     ) -> NATSMessage:
         """
         Build a command payload in the expected format.
@@ -643,17 +783,24 @@ class CommandService:
         Args:
             command_request: CommandRequest model containing command details
             machine_id: Machine ID for the command
-            run_id: Run ID for the command
+            run_id: Run ID for the command (empty string will be converted to None)
+            user_id: User ID who initiated the command
+            username: Username who initiated the command
         
         Returns:
             NATSMessage object ready for NATS transmission
         """
+        # Convert empty string to None for run_id
+        run_id_value = run_id if run_id else None
+        
         header = MessageHeader(
             message_type=MessageType.COMMAND,
             version="1.0",
             timestamp=datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ'),
+            user_id=user_id,
+            username=username,
             machine_id=machine_id,
-            run_id=run_id
+            run_id=run_id_value
         )
         
         return NATSMessage(