avtomatika 1.0b8__py3-none-any.whl → 1.0b9__py3-none-any.whl

avtomatika/engine.py

@@ -24,6 +24,7 @@ from .app_keys import (
     SCHEDULER_TASK_KEY,
     WATCHER_KEY,
     WATCHER_TASK_KEY,
+    WORKER_SERVICE_KEY,
     WS_MANAGER_KEY,
 )
 from .blueprint import StateMachineBlueprint
@@ -40,6 +41,7 @@ from .logging_config import setup_logging
 from .reputation import ReputationCalculator
 from .s3 import S3Service
 from .scheduler import Scheduler
+from .services.worker_service import WorkerService
 from .storage.base import StorageBackend
 from .telemetry import setup_telemetry
 from .utils.webhook_sender import WebhookPayload, WebhookSender
@@ -56,7 +58,7 @@ def json_dumps(obj: Any) -> str:
     return dumps(obj).decode("utf-8")
 
 
-def json_response(data: Any, **kwargs: Any) -> web.Response:
+def json_response(data, **kwargs: Any) -> web.Response:
     return web.json_response(data, dumps=json_dumps, **kwargs)
 
 
@@ -71,8 +73,13 @@ class OrchestratorEngine:
         self.ws_manager = WebSocketManager()
         self.app = web.Application(middlewares=[compression_middleware])
         self.app[ENGINE_KEY] = self
+        self.worker_service = None
         self._setup_done = False
 
+        from rxon import HttpListener
+
+        self.rxon_listener = HttpListener(self.app)
+
     def register_blueprint(self, blueprint: StateMachineBlueprint) -> None:
         if self._setup_done:
             raise RuntimeError("Cannot register blueprints after engine setup.")
@@ -142,8 +149,71 @@ class OrchestratorEngine:
                 )
                 self.history_storage = NoOpHistoryStorage()
 
+    async def handle_rxon_message(self, message_type: str, payload: Any, context: dict) -> Any:
+        """Core handler for RXON protocol messages via any listener."""
+        from rxon.security import extract_cert_identity
+
+        from .security import verify_worker_auth
+
+        request = context.get("raw_request")
+        token = context.get("token")
+        cert_identity = extract_cert_identity(request) if request else None
+
+        worker_id_hint = context.get("worker_id_hint")
+
+        if not worker_id_hint:
+            if message_type == "poll" and isinstance(payload, str):
+                worker_id_hint = payload
+            elif isinstance(payload, dict) and "worker_id" in payload:
+                worker_id_hint = payload["worker_id"]
+            elif hasattr(payload, "worker_id"):
+                worker_id_hint = payload.worker_id
+
+        try:
+            auth_worker_id = await verify_worker_auth(self.storage, self.config, token, cert_identity, worker_id_hint)
+        except PermissionError as e:
+            raise web.HTTPUnauthorized(text=str(e)) from e
+        except ValueError as e:
+            raise web.HTTPBadRequest(text=str(e)) from e
+
+        if message_type == "register":
+            return await self.worker_service.register_worker(payload)
+
+        elif message_type == "poll":
+            return await self.worker_service.get_next_task(auth_worker_id)
+
+        elif message_type == "result":
+            return await self.worker_service.process_task_result(payload, auth_worker_id)
+
+        elif message_type == "heartbeat":
+            return await self.worker_service.update_worker_heartbeat(auth_worker_id, payload)
+
+        elif message_type == "sts_token":
+            if cert_identity is None:
+                raise web.HTTPForbidden(text="Unauthorized: mTLS certificate required to issue access token.")
+            return await self.worker_service.issue_access_token(auth_worker_id)
+
+        elif message_type == "websocket":
+            ws = payload
+            await self.ws_manager.register(auth_worker_id, ws)
+            try:
+                from aiohttp import WSMsgType
+
+                async for msg in ws:
+                    if msg.type == WSMsgType.TEXT:
+                        try:
+                            data = msg.json()
+                            await self.ws_manager.handle_message(auth_worker_id, data)
+                        except Exception as e:
+                            logger.error(f"Error processing WebSocket message from {auth_worker_id}: {e}")
+                    elif msg.type == WSMsgType.ERROR:
+                        break
+            finally:
+                await self.ws_manager.unregister(auth_worker_id)
+            return None
+
     async def on_startup(self, app: web.Application) -> None:
-        # 1. Fail Fast: Check Storage Connection
+        # Fail Fast: Check Storage Connection
         if not await self.storage.ping():
             logger.critical("Failed to connect to Storage Backend (Redis). Exiting.")
             raise RuntimeError("Storage Backend is unavailable.")
@@ -208,14 +278,21 @@ class OrchestratorEngine:
         app[WS_MANAGER_KEY] = self.ws_manager
         app[S3_SERVICE_KEY] = S3Service(self.config, self.history_storage)
 
+        self.worker_service = WorkerService(self.storage, self.history_storage, self.config, self)
+        app[WORKER_SERVICE_KEY] = self.worker_service
+
         app[EXECUTOR_TASK_KEY] = create_task(app[EXECUTOR_KEY].run())
         app[WATCHER_TASK_KEY] = create_task(app[WATCHER_KEY].run())
         app[REPUTATION_CALCULATOR_TASK_KEY] = create_task(app[REPUTATION_CALCULATOR_KEY].run())
         app[HEALTH_CHECKER_TASK_KEY] = create_task(app[HEALTH_CHECKER_KEY].run())
         app[SCHEDULER_TASK_KEY] = create_task(app[SCHEDULER_KEY].run())
 
+        await self.rxon_listener.start(self.handle_rxon_message)
+
     async def on_shutdown(self, app: web.Application) -> None:
         logger.info("Shutdown sequence started.")
+        await self.rxon_listener.stop()
+
         app[EXECUTOR_KEY].stop()
         app[WATCHER_KEY].stop()
         app[REPUTATION_CALCULATOR_KEY].stop()
@@ -274,6 +351,7 @@ class OrchestratorEngine:
         blueprint_name: str,
         initial_data: dict[str, Any],
         source: str = "internal",
+        tracing_context: dict[str, str] | None = None,
     ) -> str:
         """Creates a job directly, bypassing the HTTP API layer.
         Useful for internal schedulers and triggers.
@@ -297,7 +375,7 @@ class OrchestratorEngine:
             "initial_data": initial_data,
             "state_history": {},
             "status": JOB_STATUS_PENDING,
-            "tracing_context": {},
+            "tracing_context": tracing_context or {},
             "client_config": client_config,
         }
         await self.storage.save_job_state(job_id, job_state)
@@ -374,19 +452,44 @@ class OrchestratorEngine:
 
     def run(self) -> None:
         self.setup()
+        ssl_context = None
+        if self.config.TLS_ENABLED:
+            from rxon.security import create_server_ssl_context
+
+            ssl_context = create_server_ssl_context(
+                cert_path=self.config.TLS_CERT_PATH,
+                key_path=self.config.TLS_KEY_PATH,
+                ca_path=self.config.TLS_CA_PATH,
+                require_client_cert=self.config.TLS_REQUIRE_CLIENT_CERT,
+            )
+            print(f"TLS enabled. mTLS required: {self.config.TLS_REQUIRE_CLIENT_CERT}")
+
         print(
             f"Starting OrchestratorEngine API server on {self.config.API_HOST}:{self.config.API_PORT} in blocking mode."
         )
-        web.run_app(self.app, host=self.config.API_HOST, port=self.config.API_PORT)
+        web.run_app(self.app, host=self.config.API_HOST, port=self.config.API_PORT, ssl_context=ssl_context)
 
     async def start(self):
         """Starts the orchestrator engine non-blockingly."""
         self.setup()
         self.runner = web.AppRunner(self.app)
         await self.runner.setup()
-        self.site = web.TCPSite(self.runner, self.config.API_HOST, self.config.API_PORT)
+
+        ssl_context = None
+        if self.config.TLS_ENABLED:
+            from rxon.security import create_server_ssl_context
+
+            ssl_context = create_server_ssl_context(
+                cert_path=self.config.TLS_CERT_PATH,
+                key_path=self.config.TLS_KEY_PATH,
+                ca_path=self.config.TLS_CA_PATH,
+                require_client_cert=self.config.TLS_REQUIRE_CLIENT_CERT,
+            )
+
+        self.site = web.TCPSite(self.runner, self.config.API_HOST, self.config.API_PORT, ssl_context=ssl_context)
         await self.site.start()
-        print(f"OrchestratorEngine API server running on http://{self.config.API_HOST}:{self.config.API_PORT}")
+        protocol = "https" if self.config.TLS_ENABLED else "http"
+        print(f"OrchestratorEngine API server running on {protocol}://{self.config.API_HOST}:{self.config.API_PORT}")
 
     async def stop(self):
         """Stops the orchestrator engine."""

avtomatika/executor.py

@@ -48,6 +48,16 @@ except ImportError:
     TraceContextTextMapPropagator = NoOpTraceContextTextMapPropagator()  # Instantiate the class
 
 from .app_keys import S3_SERVICE_KEY
+from .constants import (
+    JOB_STATUS_ERROR,
+    JOB_STATUS_FAILED,
+    JOB_STATUS_FINISHED,
+    JOB_STATUS_PENDING,
+    JOB_STATUS_QUARANTINED,
+    JOB_STATUS_RUNNING,
+    JOB_STATUS_WAITING_FOR_PARALLEL,
+    JOB_STATUS_WAITING_FOR_WORKER,
+)
 from .context import ActionFactory
 from .data_types import ClientConfig, JobContext
 from .history.base import HistoryStorageBase
@@ -59,7 +69,7 @@ logger = getLogger(
     __name__
 )  # Re-declare logger after potential redefinition in except block if opentelemetry was missing
 
-TERMINAL_STATES = {"finished", "failed", "error", "quarantined"}
+TERMINAL_STATES = {JOB_STATUS_FINISHED, JOB_STATUS_FAILED, JOB_STATUS_ERROR, JOB_STATUS_QUARANTINED}
 
 
 class JobExecutor:
@@ -263,7 +273,7 @@ class JobExecutor:
         # When transitioning to a new state, reset the retry counter.
         job_state["retry_count"] = 0
         job_state["current_state"] = next_state
-        job_state["status"] = "running"
+        job_state["status"] = JOB_STATUS_RUNNING
         await self.storage.save_job_state(job_id, job_state)
 
         if next_state not in TERMINAL_STATES:
@@ -280,11 +290,7 @@ class JobExecutor:
                     create_task(task_files.cleanup())
 
             await self._check_and_resume_parent(job_state)
-            # Send webhook for finished/failed jobs
-            event_type = "job_finished" if next_state == "finished" else "job_failed"
-            # Since _check_and_resume_parent is for sub-jobs, we only send webhook if it's a top-level job
-            # or if the user explicitly requested it for sub-jobs (by providing webhook_url).
-            # The current logic stores webhook_url in job_state, so we just check it.
+            event_type = "job_finished" if next_state == JOB_STATUS_FINISHED else "job_failed"
             await self.engine.send_job_webhook(job_state, event_type)
 
     async def _handle_dispatch(
@@ -313,21 +319,15 @@ class JobExecutor:
             logger.info(f"Job {job_id} is now paused, awaiting human approval.")
         else:
             logger.info(f"Job {job_id} dispatching task: {task_info}")
-
             now = monotonic()
-            # Safely get timeout, falling back to the global config if not provided in the task.
-            # This prevents TypeErrors if 'timeout_seconds' is missing.
             timeout_seconds = task_info.get("timeout_seconds") or self.engine.config.WORKER_TIMEOUT_SECONDS
             timeout_at = now + timeout_seconds
-
-            # Set status to waiting and add to watch list *before* dispatching
-            job_state["status"] = "waiting_for_worker"
+            job_state["status"] = JOB_STATUS_WAITING_FOR_WORKER
             job_state["task_dispatched_at"] = now
             job_state["current_task_info"] = task_info  # Save for retries
             job_state["current_task_transitions"] = task_info.get("transitions", {})
             await self.storage.save_job_state(job_id, job_state)
             await self.storage.add_job_to_watch(job_id, timeout_at)
-
             await self.dispatcher.dispatch(job_state, task_info)
 
     async def _handle_run_blueprint(
@@ -355,7 +355,7 @@ class JobExecutor:
             "blueprint_name": sub_blueprint_info["blueprint_name"],
             "current_state": "start",
             "initial_data": sub_blueprint_info["initial_data"],
-            "status": "pending",
+            "status": JOB_STATUS_PENDING,
             "parent_job_id": parent_job_id,
         }
         await self.storage.save_job_state(child_job_id, child_job_state)
@@ -388,7 +388,7 @@ class JobExecutor:
         branch_task_ids = [str(uuid4()) for _ in tasks_to_dispatch]
 
         # Update job state for parallel execution
-        job_state["status"] = "waiting_for_parallel_tasks"
+        job_state["status"] = JOB_STATUS_WAITING_FOR_PARALLEL
         job_state["aggregation_target"] = aggregate_into
         job_state["active_branches"] = branch_task_ids
         job_state["aggregation_results"] = {}
@@ -466,7 +466,7 @@ class JobExecutor:
             logger.critical(
                 f"Job {job_id} has failed handler execution {max_retries + 1} times. Moving to quarantine.",
             )
-            job_state["status"] = "quarantined"
+            job_state["status"] = JOB_STATUS_QUARANTINED
             job_state["error_message"] = str(error)
             await self.storage.save_job_state(job_id, job_state)
             await self.storage.quarantine_job(job_id)
@@ -499,7 +499,7 @@ class JobExecutor:
             return
 
         # Determine the outcome of the child job to select the correct transition.
-        child_outcome = "success" if child_job_state["current_state"] == "finished" else "failure"
+        child_outcome = "success" if child_job_state["current_state"] == JOB_STATUS_FINISHED else "failure"
         transitions = parent_job_state.get("current_task_transitions", {})
         next_state = transitions.get(child_outcome, "failed")
 
@@ -514,7 +514,7 @@ class JobExecutor:
 
         # Update the parent job to its new state and re-enqueue it.
         parent_job_state["current_state"] = next_state
-        parent_job_state["status"] = "running"
+        parent_job_state["status"] = JOB_STATUS_RUNNING
         await self.storage.save_job_state(parent_job_id, parent_job_state)
         await self.storage.enqueue_job(parent_job_id)
 

avtomatika/s3.py

@@ -3,13 +3,15 @@ from logging import getLogger
 from os import sep, walk
 from pathlib import Path
 from shutil import rmtree
-from typing import Any, Tuple
+from typing import Any
 
 from aiofiles import open as aiopen
 from obstore import delete_async, get_async, put_async
 from obstore import list as obstore_list
 from obstore.store import S3Store
 from orjson import dumps, loads
+from rxon.blob import calculate_config_hash, parse_uri
+from rxon.exceptions import IntegrityError
 
 from .config import Config
 from .history.base import HistoryStorageBase
@@ -56,40 +58,50 @@ class TaskFiles:
         clean_name = filename.split("/")[-1] if "://" in filename else filename.lstrip("/")
         return self.local_dir / clean_name
 
-    def _parse_s3_uri(self, uri: str) -> Tuple[str, str, bool]:
-        """
-        Parses s3://bucket/key into (bucket, key, is_directory).
-        is_directory is True if uri ends with '/'.
+    async def _download_single_file(
+        self,
+        key: str,
+        local_path: Path,
+        expected_size: int | None = None,
+        expected_hash: str | None = None,
+    ) -> dict[str, Any]:
+        """Downloads a single file safely using semaphore and streaming.
+        Returns metadata (size, etag).
         """
-        is_dir = uri.endswith("/")
-
-        if not uri.startswith("s3://"):
-            key = f"{self._s3_prefix}{uri.lstrip('/')}"
-            return self._bucket, key, is_dir
-
-        parts = uri[5:].split("/", 1)
-        bucket = parts[0]
-        key = parts[1] if len(parts) > 1 else ""
-        return bucket, key, is_dir
-
-    async def _download_single_file(self, key: str, local_path: Path) -> None:
-        """Downloads a single file safely using semaphore and streaming to avoid OOM."""
         if not local_path.parent.exists():
             await to_thread(local_path.parent.mkdir, parents=True, exist_ok=True)
 
         async with self._semaphore:
             response = await get_async(self._store, key)
+            meta = response.meta
+            file_size = meta.size
+            etag = meta.e_tag.strip('"') if meta.e_tag else None
+
+            if expected_size is not None and file_size != expected_size:
+                raise IntegrityError(f"File size mismatch for {key}: expected {expected_size}, got {file_size}")
+
+            if expected_hash is not None and etag and expected_hash != etag:
+                raise IntegrityError(f"Integrity mismatch for {key}: expected ETag {expected_hash}, got {etag}")
+
             stream = response.stream()
             async with aiopen(local_path, "wb") as f:
                 async for chunk in stream:
                     await f.write(chunk)
 
-    async def download(self, name_or_uri: str, local_name: str | None = None) -> Path:
+            return {"size": file_size, "etag": etag}
+
+    async def download(
+        self,
+        name_or_uri: str,
+        local_name: str | None = None,
+        verify_meta: dict[str, Any] | None = None,
+    ) -> Path:
         """
         Downloads a file or directory (recursively).
         If URI ends with '/', it treats it as a directory.
         """
-        bucket, key, is_dir = self._parse_s3_uri(name_or_uri)
+        bucket, key, is_dir = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
+        verify_meta = verify_meta or {}
 
         if local_name:
             target_path = self.path(local_name)
@@ -112,22 +124,42 @@ class TaskFiles:
                 tasks.append(self._download_single_file(s3_key, local_file_path))
 
             if tasks:
-                await gather(*tasks)
-
-            await self._log_event("download_dir", f"s3://{bucket}/{key}", str(target_path))
+                results = await gather(*tasks)
+                total_size = sum(r["size"] for r in results)
+                await self._log_event(
+                    "download_dir",
+                    f"s3://{bucket}/{key}",
+                    str(target_path),
+                    metadata={"total_size": total_size, "file_count": len(results)},
+                )
+            else:
+                await self._log_event(
+                    "download_dir",
+                    f"s3://{bucket}/{key}",
+                    str(target_path),
+                    metadata={"total_size": 0, "file_count": 0},
+                )
             return target_path
         else:
             logger.debug(f"Downloading s3://{bucket}/{key} -> {target_path}")
-            await self._download_single_file(key, target_path)
-            await self._log_event("download", f"s3://{bucket}/{key}", str(target_path))
+            meta = await self._download_single_file(
+                key,
+                target_path,
+                expected_size=verify_meta.get("size"),
+                expected_hash=verify_meta.get("hash"),
+            )
+            await self._log_event("download", f"s3://{bucket}/{key}", str(target_path), metadata=meta)
             return target_path
 
-    async def _upload_single_file(self, local_path: Path, s3_key: str) -> None:
-        """Uploads a single file safely using semaphore."""
+    async def _upload_single_file(self, local_path: Path, s3_key: str) -> dict[str, Any]:
+        """Uploads a single file safely using semaphore. Returns S3 metadata."""
         async with self._semaphore:
+            file_size = local_path.stat().st_size
             async with aiopen(local_path, "rb") as f:
                 content = await f.read()
-            await put_async(self._store, s3_key, content)
+            result = await put_async(self._store, s3_key, content)
+            etag = result.e_tag.strip('"') if result.e_tag else None
+            return {"size": file_size, "etag": etag}
 
     async def upload(self, local_name: str, remote_name: str | None = None) -> str:
         """
@@ -158,26 +190,30 @@ class TaskFiles:
 
             tasks = [self._upload_single_file(lp, k) for lp, k in files_map]
             if tasks:
-                await gather(*tasks)
+                results = await gather(*tasks)
+                total_size = sum(r["size"] for r in results)
+                metadata = {"total_size": total_size, "file_count": len(results)}
+            else:
+                metadata = {"total_size": 0, "file_count": 0}
 
             uri = f"s3://{self._bucket}/{target_prefix}"
-            await self._log_event("upload_dir", uri, str(local_path))
+            await self._log_event("upload_dir", uri, str(local_path), metadata=metadata)
             return uri
 
         elif local_path.exists():
             target_key = f"{self._s3_prefix}{(remote_name or local_name).lstrip('/')}"
             logger.debug(f"Uploading {local_path} -> s3://{self._bucket}/{target_key}")
 
-            await self._upload_single_file(local_path, target_key)
+            meta = await self._upload_single_file(local_path, target_key)
 
             uri = f"s3://{self._bucket}/{target_key}"
-            await self._log_event("upload", uri, str(local_path))
+            await self._log_event("upload", uri, str(local_path), metadata=meta)
             return uri
         else:
             raise FileNotFoundError(f"Local file/dir not found: {local_path}")
 
     async def read_text(self, name_or_uri: str) -> str:
-        bucket, key, _ = self._parse_s3_uri(name_or_uri)
+        bucket, key, _ = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
         filename = key.split("/")[-1]
         local_path = self.path(filename)
 
@@ -188,7 +224,7 @@ class TaskFiles:
             return await f.read()
 
     async def read_json(self, name_or_uri: str) -> Any:
-        bucket, key, _ = self._parse_s3_uri(name_or_uri)
+        bucket, key, _ = parse_uri(name_or_uri, self._bucket, self._s3_prefix)
         filename = key.split("/")[-1]
         local_path = self.path(filename)
 
@@ -235,21 +271,31 @@ class TaskFiles:
         if self.local_dir.exists():
             await to_thread(rmtree, self.local_dir)
 
-    async def _log_event(self, operation: str, file_uri: str, local_path: str) -> None:
+    async def _log_event(
+        self,
+        operation: str,
+        file_uri: str,
+        local_path: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> None:
         if not self._history:
             return
 
         try:
+            context_snapshot = {
+                "operation": operation,
+                "s3_uri": file_uri,
+                "local_path": str(local_path),
+            }
+            if metadata:
+                context_snapshot.update(metadata)
+
             await self._history.log_job_event(
                 {
                     "job_id": self._job_id,
                     "event_type": "s3_operation",
                     "state": "running",
-                    "context_snapshot": {
-                        "operation": operation,
-                        "s3_uri": file_uri,
-                        "local_path": str(local_path),
-                    },
+                    "context_snapshot": context_snapshot,
                 }
             )
         except Exception as e:
@@ -306,6 +352,16 @@ class S3Service:
             logger.error(f"Failed to initialize S3 Store: {e}")
             self._enabled = False
 
+    def get_config_hash(self) -> str | None:
+        """Returns a hash of the current S3 configuration for consistency checks."""
+        if not self._enabled:
+            return None
+        return calculate_config_hash(
+            self.config.S3_ENDPOINT_URL,
+            self.config.S3_ACCESS_KEY,
+            self.config.S3_DEFAULT_BUCKET,
+        )
+
     def get_task_files(self, job_id: str) -> TaskFiles | None:
         if not self._enabled or not self._store or not self._semaphore:
             return None