Merge pull request #10 from codeSamuraii/copilot/add-worker-logging-feature

Harden input validation, security, and packaging for first open source release

Author: codeSamuraii

pyfuse/core/__init__.py

@@ -2,3 +2,14 @@
 from pyfuse.core.errors import Error, RemoteError, WorkerError, DependencyError
 from pyfuse.core.models import ImportInfo, FunctionNode
 from pyfuse.core.version import _VERSION
+
+__all__ = [
+    "Task",
+    "resolve_args",
+    "Error",
+    "RemoteError",
+    "WorkerError",
+    "DependencyError",
+    "ImportInfo",
+    "FunctionNode",
+]

pyfuse/core/errors.py

@@ -71,4 +71,5 @@ def _pyfuse_excepthook(
     _original_excepthook(exc_type, exc_value, exc_tb)
 
 
-sys.excepthook = _pyfuse_excepthook
+if sys.excepthook is not _pyfuse_excepthook:
+    sys.excepthook = _pyfuse_excepthook

pyfuse/core/models.py

@@ -23,7 +23,13 @@ def to_dict(self) -> dict[str, str]:
 
     @classmethod
     def from_dict(cls, data: dict[str, str]) -> Self:
-        """Deserialize from a plain dict."""
+        """Deserialize from a plain dict.
+
+        Raises
+        ------
+        KeyError
+            If required fields (``statement``, ``bound_name``) are missing.
+        """
         return cls(
             statement=data["statement"],
             bound_name=data["bound_name"],
@@ -134,7 +140,18 @@ def content_hash(self) -> str:
 
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> Self:
-        """Deserialize from a plain dict."""
+        """Deserialize from a plain dict.
+
+        Raises
+        ------
+        KeyError
+            If required fields (``qualified_name``, ``name``, ``module``,
+            ``source``, ``imports``, ``dependencies``) are missing.
+        """
+        required = ("qualified_name", "name", "module", "source", "imports", "dependencies")
+        missing = [k for k in required if k not in data]
+        if missing:
+            raise KeyError(f"FunctionNode missing required fields: {', '.join(missing)}")
         return cls(
             qualified_name=data["qualified_name"],
             name=data["name"],

pyfuse/core/task.py

@@ -166,7 +166,7 @@ def from_json(
         sig = data.pop("signature", None)
 
         if signing_key is not None:
-            if sig is None:
+            if not sig:
                 raise SignatureError(
                     "Task is unsigned but signing is enabled — "
                     "rejecting unauthenticated task"
@@ -185,5 +185,5 @@ def from_json(
             timeout=data.get("timeout"),
             retries=data.get("retries", 0),
             retry_delay=data.get("retry_delay", 1.0),
-            signature=sig,
+            signature=sig or None,  # normalise empty string to None
         )

pyfuse/graph/__init__.py

@@ -1,3 +1,5 @@
 from pyfuse.graph.graph import Graph
 from pyfuse.graph.store import Store, MergeResult
 from pyfuse.graph.decorator import trace
+
+__all__ = ["Graph", "Store", "MergeResult", "trace"]

pyfuse/graph/decorator.py

@@ -55,6 +55,13 @@ def _apply_trace(
     retries: int = 0,
     retry_delay: float = 1.0,
 ) -> TracedFunction[_P, _R]:
+    if timeout is not None and timeout <= 0:
+        raise ValueError(f"timeout must be positive, got {timeout}")
+    if retries < 0:
+        raise ValueError(f"retries must be non-negative, got {retries}")
+    if retry_delay < 0:
+        raise ValueError(f"retry_delay must be non-negative, got {retry_delay}")
+
     logger.debug("@trace applied to %s", func.__qualname__)
     graph = Graph.default()
     graph.register(func)

pyfuse/worker/__init__.py

@@ -2,3 +2,16 @@
 from pyfuse.worker.remote import serve, connect, disconnect, submit_remote
 from pyfuse.worker.result import Result, ResultEnvelope
 from pyfuse.worker.worker import Worker, execute
+
+__all__ = [
+    "install_package_as",
+    "ensure_dependencies",
+    "serve",
+    "connect",
+    "disconnect",
+    "submit_remote",
+    "Result",
+    "ResultEnvelope",
+    "Worker",
+    "execute",
+]

pyfuse/worker/backends/__init__.py

@@ -1 +1,3 @@
 from pyfuse.worker.backends.base import Backend
+
+__all__ = ["Backend"]

pyfuse/worker/deps.py

@@ -96,8 +96,10 @@ async def _pip_install(package: str, extra_args: list[str]) -> tuple[int, str, s
         stderr=asyncio.subprocess.PIPE,
     )
     stdout_bytes, stderr_bytes = await proc.communicate()
+    # returncode is None if the process hasn't terminated; treat as failure
+    returncode = proc.returncode if proc.returncode is not None else 1
     return (
-        proc.returncode or 0,
+        returncode,
         stdout_bytes.decode() if stdout_bytes else "",
         stderr_bytes.decode() if stderr_bytes else "",
     )

pyfuse/worker/remote.py

@@ -220,14 +220,14 @@ async def _heartbeat_loop(
         try:
             await backend.send_heartbeat(task_id)
         except Exception:
-            pass  # best-effort
+            logger.debug("Heartbeat send failed for task %s", task_id, exc_info=True)
         if exec_task is not None:
             try:
                 if await backend.is_cancelled(task_id):
                     exec_task.cancel()
                     return
             except Exception:
-                pass
+                logger.debug("Cancellation check failed for task %s", task_id, exc_info=True)
         try:
             await asyncio.wait_for(cancel_event.wait(), timeout=_HEARTBEAT_INTERVAL)
         except asyncio.TimeoutError: