penguiflow 2.0.0__py3-none-any.whl → 2.2.0__py3-none-any.whl

penguiflow/core.py

@@ -9,17 +9,20 @@ from __future__ import annotations
 import asyncio
 import logging
 import time
+import warnings
 from collections import deque
 from collections.abc import Awaitable, Callable, Mapping, Sequence
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 
+from .bus import BusEnvelope, MessageBus
 from .errors import FlowError, FlowErrorCode
 from .metrics import FlowEvent
 from .middlewares import Middleware
 from .node import Node, NodePolicy
 from .registry import ModelRegistry
+from .state import RemoteBinding, StateStore, StoredEvent
 from .types import WM, FinalAnswer, Message, StreamChunk
 
 logger = logging.getLogger("penguiflow.core")
@@ -143,8 +146,7 @@ class Context:
         if self._runtime is None:
             raise RuntimeError("Context is not attached to a running flow")
         for floe in self._resolve_targets(to, self._outgoing):
-            self._runtime._on_message_enqueued(msg)
-            await floe.queue.put(msg)
+            await self._runtime._send_to_floe(floe, msg)
 
     def emit_nowait(
         self, msg: Any, to: Node | Endpoint | Sequence[Node | Endpoint] | None = None
@@ -152,8 +154,7 @@ class Context:
         if self._runtime is None:
             raise RuntimeError("Context is not attached to a running flow")
         for floe in self._resolve_targets(to, self._outgoing):
-            self._runtime._on_message_enqueued(msg)
-            floe.queue.put_nowait(msg)
+            self._runtime._send_to_floe_nowait(floe, msg)
 
     async def emit_chunk(
         self,
@@ -301,6 +302,8 @@ class PenguiFlow:
         allow_cycles: bool = False,
         middlewares: Sequence[Middleware] | None = None,
         emit_errors_to_rookery: bool = False,
+        state_store: StateStore | None = None,
+        message_bus: MessageBus | None = None,
     ) -> None:
         self._queue_maxsize = queue_maxsize
         self._allow_cycles = allow_cycles
@@ -314,10 +317,14 @@ class PenguiFlow:
         self._middlewares: list[Middleware] = list(middlewares or [])
         self._trace_counts: dict[str, int] = {}
         self._trace_events: dict[str, asyncio.Event] = {}
-        self._trace_invocations: dict[str, set[asyncio.Future[Any]]] = {}
+        self._trace_invocations: dict[str, set[asyncio.Task[Any]]] = {}
+        self._external_tasks: dict[str, set[asyncio.Future[Any]]] = {}
         self._trace_capacity_waiters: dict[str, list[asyncio.Event]] = {}
         self._latest_wm_hops: dict[str, int] = {}
         self._emit_errors_to_rookery = emit_errors_to_rookery
+        self._state_store = state_store
+        self._message_bus = message_bus
+        self._bus_tasks: set[asyncio.Task[None]] = set()
 
         self._build_graph(adjacencies)
 
@@ -487,6 +494,29 @@ class PenguiFlow:
             task.cancel()
         await asyncio.gather(*self._tasks, return_exceptions=True)
         self._tasks.clear()
+        if self._trace_invocations:
+            pending: list[asyncio.Task[Any]] = []
+            for invocation_tasks in self._trace_invocations.values():
+                for task in invocation_tasks:
+                    if not task.done():
+                        task.cancel()
+                    pending.append(task)
+            if pending:
+                await asyncio.gather(*pending, return_exceptions=True)
+            self._trace_invocations.clear()
+        if self._external_tasks:
+            pending_ext: list[asyncio.Future[Any]] = []
+            for external_tasks in self._external_tasks.values():
+                for external_task in external_tasks:
+                    if not external_task.done():
+                        external_task.cancel()
+                    pending_ext.append(external_task)
+            if pending_ext:
+                await asyncio.gather(*pending_ext, return_exceptions=True)
+            self._external_tasks.clear()
+        if self._bus_tasks:
+            await asyncio.gather(*self._bus_tasks, return_exceptions=True)
+            self._bus_tasks.clear()
         self._trace_counts.clear()
         self._trace_events.clear()
         self._trace_invocations.clear()
@@ -547,6 +577,84 @@ class PenguiFlow:
         await self._finalize_message(result)
         return result
 
+    async def load_history(self, trace_id: str) -> Sequence[StoredEvent]:
+        """Return the persisted history for ``trace_id`` from the state store."""
+
+        if self._state_store is None:
+            raise RuntimeError("PenguiFlow was created without a state_store")
+        return await self._state_store.load_history(trace_id)
+
+    def ensure_trace_event(self, trace_id: str) -> asyncio.Event:
+        """Return (and create if needed) the cancellation event for ``trace_id``."""
+
+        return self._trace_events.setdefault(trace_id, asyncio.Event())
+
+    def register_external_task(self, trace_id: str, task: asyncio.Future[Any]) -> None:
+        """Track an externally created task for cancellation bookkeeping."""
+
+        if trace_id is None:
+            return
+        tasks = self._external_tasks.get(trace_id)
+        if tasks is None:
+            tasks = set[asyncio.Future[Any]]()
+            self._external_tasks[trace_id] = tasks
+        tasks.add(task)
+
+        def _cleanup(finished: asyncio.Future[Any]) -> None:
+            remaining = self._external_tasks.get(trace_id)
+            if remaining is None:
+                return
+            remaining.discard(finished)
+            if not remaining:
+                self._external_tasks.pop(trace_id, None)
+
+        task.add_done_callback(_cleanup)
+
+    async def save_remote_binding(self, binding: RemoteBinding) -> None:
+        """Persist a remote binding if a state store is configured."""
+
+        if self._state_store is None:
+            return
+        try:
+            await self._state_store.save_remote_binding(binding)
+        except Exception as exc:  # pragma: no cover - defensive logging
+            logger.exception(
+                "state_store_binding_failed",
+                extra={
+                    "event": "state_store_binding_failed",
+                    "trace_id": binding.trace_id,
+                    "context_id": binding.context_id,
+                    "task_id": binding.task_id,
+                    "agent_url": binding.agent_url,
+                    "exception": repr(exc),
+                },
+            )
+
+    async def record_remote_event(
+        self,
+        *,
+        event: str,
+        node: Node,
+        context: Context,
+        trace_id: str | None,
+        latency_ms: float | None,
+        level: int = logging.INFO,
+        extra: Mapping[str, Any] | None = None,
+    ) -> None:
+        """Emit a structured :class:`FlowEvent` for remote transport activity."""
+
+        payload = dict(extra or {})
+        await self._emit_event(
+            event=event,
+            node=node,
+            context=context,
+            trace_id=trace_id,
+            attempt=0,
+            latency_ms=latency_ms,
+            level=level,
+            extra=payload,
+        )
+
     async def _execute_with_reliability(
         self,
         node: Node,
@@ -579,6 +687,21 @@ class PenguiFlow:
                     trace_id,
                 )
 
+                if (
+                    result is not None
+                    and self._expects_message_output(node)
+                    and not isinstance(result, Message)
+                ):
+                    node_name = node.name or node.node_id
+                    warning_msg = (
+                        "Node "
+                        f"'{node_name}' is registered for Message -> Message outputs "
+                        f"but returned {type(result).__name__}. "
+                        "Return a penguiflow.types.Message to preserve headers, "
+                        "trace_id, and meta."
+                    )
+                    warnings.warn(warning_msg, RuntimeWarning, stacklevel=2)
+
                 if result is not None:
                     (
                         destination,
@@ -805,16 +928,19 @@ class PenguiFlow:
         return await self._await_invocation(node, invocation, trace_id, timeout)
 
     def _register_invocation_task(
-        self, trace_id: str, task: asyncio.Future[Any]
+        self, trace_id: str, task: asyncio.Task[Any]
     ) -> None:
-        tasks = self._trace_invocations.setdefault(trace_id, set())
+        tasks = self._trace_invocations.get(trace_id)
+        if tasks is None:
+            tasks = set[asyncio.Task[Any]]()
+            self._trace_invocations[trace_id] = tasks
         tasks.add(task)
 
         def _cleanup(finished: asyncio.Future[Any]) -> None:
             remaining = self._trace_invocations.get(trace_id)
             if remaining is None:
                 return
-            remaining.discard(finished)
+            remaining.discard(cast(asyncio.Task[Any], finished))
             if not remaining:
                 self._trace_invocations.pop(trace_id, None)
 
@@ -827,7 +953,7 @@ class PenguiFlow:
         trace_id: str,
         timeout: float | None,
     ) -> Any:
-        invocation_task = asyncio.ensure_future(invocation)
+        invocation_task = cast(asyncio.Task[Any], asyncio.ensure_future(invocation))
         self._register_invocation_task(trace_id, invocation_task)
 
         cancel_event = self._trace_events.get(trace_id)
@@ -904,6 +1030,89 @@ class PenguiFlow:
         self._trace_counts[trace_id] = self._trace_counts.get(trace_id, 0) + 1
         self._trace_events.setdefault(trace_id, asyncio.Event())
 
+    def _node_label(self, node: Node | Endpoint | None) -> str | None:
+        if node is None:
+            return None
+        name = getattr(node, "name", None)
+        if name:
+            return name
+        return getattr(node, "node_id", None)
+
+    def _build_bus_envelope(
+        self,
+        source: Node | Endpoint | None,
+        target: Node | Endpoint | None,
+        message: Any,
+    ) -> BusEnvelope:
+        source_name = self._node_label(source)
+        target_name = self._node_label(target)
+        edge = f"{source_name or '*'}->{target_name or '*'}"
+        headers: Mapping[str, Any] | None = None
+        meta: Mapping[str, Any] | None = None
+        if isinstance(message, Message):
+            headers = message.headers.model_dump()
+            meta = dict(message.meta)
+        return BusEnvelope(
+            edge=edge,
+            source=source_name,
+            target=target_name,
+            trace_id=self._get_trace_id(message),
+            payload=message,
+            headers=headers,
+            meta=meta,
+        )
+
+    async def _publish_to_bus(
+        self,
+        source: Node | Endpoint | None,
+        target: Node | Endpoint | None,
+        message: Any,
+    ) -> None:
+        if self._message_bus is None:
+            return
+        envelope = self._build_bus_envelope(source, target, message)
+        try:
+            await self._message_bus.publish(envelope)
+        except Exception as exc:
+            logger.exception(
+                "message_bus_publish_failed",
+                extra={
+                    "event": "message_bus_publish_failed",
+                    "edge": envelope.edge,
+                    "trace_id": envelope.trace_id,
+                    "exception": repr(exc),
+                },
+            )
+
+    def _schedule_bus_publish(
+        self,
+        source: Node | Endpoint | None,
+        target: Node | Endpoint | None,
+        message: Any,
+    ) -> None:
+        if self._message_bus is None:
+            return
+        loop = asyncio.get_running_loop()
+        task = loop.create_task(self._publish_to_bus(source, target, message))
+        self._bus_tasks.add(task)
+
+        def _cleanup(done: asyncio.Task[None]) -> None:
+            self._bus_tasks.discard(done)
+
+        task.add_done_callback(_cleanup)
+
+    async def _send_to_floe(self, floe: Floe, message: Any) -> None:
+        self._on_message_enqueued(message)
+        if self._message_bus is not None:
+            await self._publish_to_bus(floe.source, floe.target, message)
+        await floe.queue.put(message)
+
+    def _send_to_floe_nowait(self, floe: Floe, message: Any) -> None:
+        self._on_message_enqueued(message)
+        if self._message_bus is not None:
+            self._schedule_bus_publish(floe.source, floe.target, message)
+        floe.queue.put_nowait(message)
+
     async def _finalize_message(self, message: Any) -> None:
         trace_id = self._get_trace_id(message)
         if trace_id is None:
@@ -1017,6 +1226,29 @@ class PenguiFlow:
         for waiter in waiters:
             waiter.set()
 
+    def _expects_message_output(self, node: Node) -> bool:
+        registry = self._registry
+        if registry is None:
+            return False
+
+        models = getattr(registry, "models", None)
+        if models is None:
+            return False
+
+        node_name = node.name
+        if not node_name:
+            return False
+
+        try:
+            _in_model, out_model = models(node_name)
+        except Exception:  # pragma: no cover - registry without entry
+            return False
+
+        try:
+            return issubclass(out_model, Message)
+        except TypeError:
+            return False
+
     def _controller_postprocess(
         self,
         node: Node,
@@ -1107,11 +1339,12 @@ class PenguiFlow:
         if floe is None and incoming:
             floe = next(iter(incoming.values()))
 
-        self._on_message_enqueued(message)
-
         if floe is not None:
-            await floe.queue.put(message)
+            await self._send_to_floe(floe, message)
         else:
+            self._on_message_enqueued(message)
+            if self._message_bus is not None:
+                await self._publish_to_bus(source, ROOKERY, message)
             buffer = rookery_context._buffer
             buffer.append(message)
 
@@ -1167,6 +1400,21 @@ class PenguiFlow:
 
         logger.log(level, event, extra=event_obj.to_payload())
 
+        if self._state_store is not None:
+            stored_event = StoredEvent.from_flow_event(event_obj)
+            try:
+                await self._state_store.save_event(stored_event)
+            except Exception as exc:
+                logger.exception(
+                    "state_store_save_failed",
+                    extra={
+                        "event": "state_store_save_failed",
+                        "trace_id": stored_event.trace_id,
+                        "kind": stored_event.kind,
+                        "exception": repr(exc),
+                    },
+                )
+
         for middleware in list(self._middlewares):
             try:
                 await middleware(event_obj)

penguiflow/debug.py

@@ -0,0 +1,30 @@
+"""Developer-facing debugging helpers for PenguiFlow."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import Any
+
+from .metrics import FlowEvent
+
+
+def format_flow_event(event: FlowEvent) -> dict[str, Any]:
+    """Return a structured payload ready for logging.
+
+    The returned dictionary mirrors :meth:`FlowEvent.to_payload` and flattens any
+    embedded ``FlowError`` payload so that log aggregators can index the error
+    metadata (``flow_error_code``, ``flow_error_message``, ...).
+    """
+
+    payload = dict(event.to_payload())
+    error_payload: Mapping[str, Any] | None = event.error_payload
+    if error_payload is not None:
+        # Preserve the original payload for downstream consumers.
+        payload["flow_error"] = dict(error_payload)
+        for key, value in error_payload.items():
+            payload[f"flow_error_{key}"] = value
+    return payload
+
+
+__all__ = ["format_flow_event"]
+

penguiflow/metrics.py

@@ -31,6 +31,15 @@ class FlowEvent:
     def __post_init__(self) -> None:
         object.__setattr__(self, "extra", MappingProxyType(dict(self.extra)))
 
+    @property
+    def error_payload(self) -> Mapping[str, Any] | None:
+        """Return the structured ``FlowError`` payload if present."""
+
+        raw_payload = self.extra.get("flow_error")
+        if isinstance(raw_payload, Mapping):
+            return MappingProxyType(dict(raw_payload))
+        return None
+
     @property
     def queue_depth(self) -> int:
         """Return the combined depth of incoming and outgoing queues."""

penguiflow/middlewares.py

@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+import logging
+from collections.abc import Callable
 from typing import Protocol
 
 from .metrics import FlowEvent
 
+LatencyCallback = Callable[[str, float, FlowEvent], None]
+
 
 class Middleware(Protocol):
     """Base middleware signature receiving :class:`FlowEvent` objects."""
@@ -13,4 +17,71 @@ class Middleware(Protocol):
     async def __call__(self, event: FlowEvent) -> None: ...
 
 
-__all__ = ["Middleware", "FlowEvent"]
+def log_flow_events(
+    logger: logging.Logger | None = None,
+    *,
+    start_level: int = logging.INFO,
+    success_level: int = logging.INFO,
+    error_level: int = logging.ERROR,
+    latency_callback: LatencyCallback | None = None,
+) -> Middleware:
+    """Return middleware that emits structured node lifecycle logs.
+
+    Parameters
+    ----------
+    logger:
+        Optional :class:`logging.Logger` instance. When omitted a logger named
+        ``"penguiflow.flow"`` is used.
+    start_level, success_level, error_level:
+        Logging levels for ``node_start``, ``node_success``, and
+        ``node_error`` events respectively.
+    latency_callback:
+        Optional callable invoked with ``(event_type, latency_ms, event)`` for
+        ``node_success`` and ``node_error`` events. Use this hook to connect the
+        middleware to histogram-based metrics backends without
+        re-implementing timing logic.
+    """
+
+    log = logger or logging.getLogger("penguiflow.flow")
+
+    async def _middleware(event: FlowEvent) -> None:
+        if event.event_type not in {"node_start", "node_success", "node_error"}:
+            return
+
+        payload = event.to_payload()
+        log_level = start_level
+
+        if event.event_type == "node_start":
+            log_level = start_level
+        elif event.event_type == "node_success":
+            log_level = success_level
+        else:
+            log_level = error_level
+            if event.error_payload is not None:
+                payload = dict(payload)
+                payload["error_payload"] = dict(event.error_payload)
+
+        log.log(log_level, event.event_type, extra=payload)
+
+        if (
+            latency_callback is not None
+            and event.event_type in {"node_success", "node_error"}
+            and event.latency_ms is not None
+        ):
+            try:
+                latency_callback(event.event_type, float(event.latency_ms), event)
+            except Exception:
+                log.exception(
+                    "log_flow_events_latency_callback_error",
+                    extra={
+                        "event": "log_flow_events_latency_callback_error",
+                        "node_name": event.node_name,
+                        "node_id": event.node_id,
+                        "trace_id": event.trace_id,
+                    },
+                )
+
+    return _middleware
+
+
+__all__ = ["Middleware", "FlowEvent", "log_flow_events", "LatencyCallback"]

penguiflow/registry.py

@@ -15,6 +15,8 @@ ModelT = TypeVar("ModelT", bound=BaseModel)
 class RegistryEntry:
     in_adapter: TypeAdapter[Any]
     out_adapter: TypeAdapter[Any]
+    in_model: type[BaseModel]
+    out_model: type[BaseModel]
 
 
 class ModelRegistry:
@@ -36,6 +38,8 @@ class ModelRegistry:
         self._entries[node_name] = RegistryEntry(
             TypeAdapter(in_model),
             TypeAdapter(out_model),
+            in_model,
+            out_model,
         )
 
     def adapters(self, node_name: str) -> tuple[TypeAdapter[Any], TypeAdapter[Any]]:
@@ -45,5 +49,22 @@ class ModelRegistry:
             raise KeyError(f"Node '{node_name}' not registered") from exc
         return entry.in_adapter, entry.out_adapter
 
+    def models(
+        self, node_name: str
+    ) -> tuple[type[BaseModel], type[BaseModel]]:
+        """Return the registered models for ``node_name``.
+
+        Raises
+        ------
+        KeyError
+            If the node has not been registered.
+        """
+
+        try:
+            entry = self._entries[node_name]
+        except KeyError as exc:
+            raise KeyError(f"Node '{node_name}' not registered") from exc
+        return entry.in_model, entry.out_model
+
 
 __all__ = ["ModelRegistry"]