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

penguiflow/remote.py

@@ -0,0 +1,486 @@
+"""Remote transport protocol and helper node for PenguiFlow."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import time
+from collections.abc import AsyncIterator, Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Protocol
+
+from pydantic import BaseModel
+
+from .core import TraceCancelled
+from .node import Node, NodePolicy
+from .state import RemoteBinding
+from .types import Message
+
+if TYPE_CHECKING:  # pragma: no cover - import for typing only
+    from .core import Context, PenguiFlow
+
+
+@dataclass(slots=True)
+class RemoteCallRequest:
+    """Input to :class:`RemoteTransport` implementations."""
+
+    message: Message
+    skill: str
+    agent_url: str
+    agent_card: Mapping[str, Any] | None = None
+    metadata: Mapping[str, Any] | None = None
+    timeout_s: float | None = None
+
+
+@dataclass(slots=True)
+class RemoteCallResult:
+    """Return value for :meth:`RemoteTransport.send`."""
+
+    result: Any
+    context_id: str | None = None
+    task_id: str | None = None
+    agent_url: str | None = None
+    meta: Mapping[str, Any] | None = None
+
+
+@dataclass(slots=True)
+class RemoteStreamEvent:
+    """Streaming event yielded by :meth:`RemoteTransport.stream`."""
+
+    text: str | None = None
+    done: bool = False
+    meta: Mapping[str, Any] | None = None
+    context_id: str | None = None
+    task_id: str | None = None
+    agent_url: str | None = None
+    result: Any | None = None
+
+
+class RemoteTransport(Protocol):
+    """Protocol describing the minimal remote invocation surface."""
+
+    async def send(self, request: RemoteCallRequest) -> RemoteCallResult:
+        """Perform a unary remote call."""
+
+    def stream(self, request: RemoteCallRequest) -> AsyncIterator[RemoteStreamEvent]:
+        """Perform a remote call that yields streaming events."""
+
+    async def cancel(self, *, agent_url: str, task_id: str) -> None:
+        """Cancel a remote task identified by ``task_id`` at ``agent_url``."""
+
+
+def _json_default(value: Any) -> Any:
+    if isinstance(value, BaseModel):
+        return value.model_dump(mode="json")
+    if isinstance(value, bytes):
+        return value.decode("utf-8", errors="replace")
+    return repr(value)
+
+
+def _estimate_bytes(value: Any) -> int | None:
+    """Best-effort size estimation for observability metrics."""
+
+    if value is None:
+        return None
+    try:
+        if isinstance(value, BaseModel):
+            payload = value.model_dump(mode="json")
+        else:
+            payload = value
+        encoded = json.dumps(payload, default=_json_default).encode("utf-8")
+    except Exception:
+        try:
+            encoded = str(value).encode("utf-8")
+        except Exception:
+            return None
+    return len(encoded)
+
+
+def _text_bytes(text: str | None) -> int:
+    if text is None:
+        return 0
+    return len(text.encode("utf-8"))
+
+
+def _merge_remote_extra(
+    base: Mapping[str, Any],
+    *,
+    agent_url: str | None,
+    context_id: str | None,
+    task_id: str | None,
+    additional: Mapping[str, Any] | None = None,
+) -> dict[str, Any]:
+    extra = dict(base)
+    if agent_url is not None:
+        extra["remote_agent_url"] = agent_url
+    if context_id is not None:
+        extra["remote_context_id"] = context_id
+    if task_id is not None:
+        extra["remote_task_id"] = task_id
+    if additional:
+        for key, value in additional.items():
+            if value is not None:
+                extra[key] = value
+    return extra
+
+
+def RemoteNode(
+    *,
+    transport: RemoteTransport,
+    skill: str,
+    agent_url: str,
+    name: str,
+    agent_card: Mapping[str, Any] | None = None,
+    policy: NodePolicy | None = None,
+    streaming: bool = False,
+    record_binding: bool = True,
+) -> Node:
+    """Create a node that proxies work to a remote agent via ``transport``."""
+
+    node_policy = policy or NodePolicy()
+
+    async def _record_binding(
+        *,
+        runtime: PenguiFlow,
+        context: Context,
+        node_owner: Node,
+        trace_id: str,
+        context_id: str | None,
+        task_id: str | None,
+        agent_url_override: str | None,
+        base_extra: Mapping[str, Any],
+    ) -> tuple[asyncio.Task[None], asyncio.Event] | None:
+        if context_id is None or task_id is None:
+            return None
+
+        agent_ref = agent_url_override or agent_url
+
+        if record_binding:
+            binding = RemoteBinding(
+                trace_id=trace_id,
+                context_id=context_id,
+                task_id=task_id,
+                agent_url=agent_ref,
+            )
+            await runtime.save_remote_binding(binding)
+
+        cancel_event = runtime.ensure_trace_event(trace_id)
+
+        async def _issue_cancel(reason: str) -> None:
+            start_cancel = time.perf_counter()
+            try:
+                await transport.cancel(agent_url=agent_ref, task_id=task_id)
+            except Exception as exc:  # pragma: no cover - defensive logging
+                latency = (time.perf_counter() - start_cancel) * 1000
+                extra = _merge_remote_extra(
+                    base_extra,
+                    agent_url=agent_ref,
+                    context_id=context_id,
+                    task_id=task_id,
+                    additional={
+                        "remote_cancel_reason": reason,
+                        "remote_error": repr(exc),
+                        "remote_status": "cancel_error",
+                    },
+                )
+                await runtime.record_remote_event(
+                    event="remote_cancel_error",
+                    node=node_owner,
+                    context=context,
+                    trace_id=trace_id,
+                    latency_ms=latency,
+                    level=logging.ERROR,
+                    extra=extra,
+                )
+                return
+
+            latency = (time.perf_counter() - start_cancel) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=agent_ref,
+                context_id=context_id,
+                task_id=task_id,
+                additional={
+                    "remote_cancel_reason": reason,
+                    "remote_status": "cancelled",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_cancelled",
+                node=node_owner,
+                context=context,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.INFO,
+                extra=extra,
+            )
+
+        if cancel_event.is_set():
+            await _issue_cancel("pre_cancelled")
+            raise TraceCancelled(trace_id)
+
+        async def _mirror_cancel() -> None:
+            try:
+                await cancel_event.wait()
+            except asyncio.CancelledError:
+                return
+            await _issue_cancel("trace_cancel")
+
+        cancel_task = asyncio.create_task(_mirror_cancel())
+        runtime.register_external_task(trace_id, cancel_task)
+        return cancel_task, cancel_event
+
+    async def _remote_impl(message: Message, ctx: Context) -> Any:
+        if not isinstance(message, Message):
+            raise TypeError("Remote nodes require penguiflow.types.Message inputs")
+
+        runtime = ctx.runtime
+        if runtime is None:
+            raise RuntimeError("Context is not bound to a running PenguiFlow")
+
+        owner = ctx.owner
+        if not isinstance(owner, Node):  # pragma: no cover - defensive safety
+            raise RuntimeError("Remote context owner must be a Node")
+
+        trace_id = message.trace_id
+        cancel_task: asyncio.Task[None] | None = None
+        cancel_event: asyncio.Event | None = None
+        binding_registered = False
+
+        remote_context_id: str | None = None
+        remote_task_id: str | None = None
+        remote_agent_url_final = agent_url
+        response_bytes = 0
+        stream_events = 0
+
+        base_extra: dict[str, Any] = {
+            "remote_skill": skill,
+            "remote_transport": type(transport).__name__,
+            "remote_streaming": streaming,
+        }
+        request_bytes = _estimate_bytes(message)
+        if request_bytes is not None:
+            base_extra["remote_request_bytes"] = request_bytes
+
+        request = RemoteCallRequest(
+            message=message,
+            skill=skill,
+            agent_url=agent_url,
+            agent_card=agent_card,
+            metadata=message.meta,
+            timeout_s=node_policy.timeout_s,
+        )
+
+        async def _ensure_binding(
+            *,
+            context_id: str | None,
+            task_id: str | None,
+            agent_url_override: str | None,
+        ) -> None:
+            nonlocal cancel_task, cancel_event, binding_registered
+            nonlocal remote_context_id, remote_task_id, remote_agent_url_final
+            if context_id is not None:
+                remote_context_id = context_id
+            if task_id is not None:
+                remote_task_id = task_id
+            if agent_url_override is not None:
+                remote_agent_url_final = agent_url_override
+            if binding_registered:
+                return
+            if context_id is None or task_id is None:
+                return
+            record = await _record_binding(
+                runtime=runtime,
+                context=ctx,
+                node_owner=owner,
+                trace_id=trace_id,
+                context_id=context_id,
+                task_id=task_id,
+                agent_url_override=agent_url_override,
+                base_extra=base_extra,
+            )
+            if record is None:
+                return
+            cancel_task, cancel_event = record
+            binding_registered = True
+
+        async def _cleanup_cancel_task() -> None:
+            if cancel_task is not None:
+                try:
+                    if cancel_event is not None and cancel_event.is_set():
+                        await cancel_task
+                        return
+                    if not cancel_task.done():
+                        cancel_task.cancel()
+                    await cancel_task
+                except BaseException:  # pragma: no cover - cleanup guard
+                    pass
+
+        async def _run_stream() -> Any | None:
+            nonlocal response_bytes, stream_events, remote_agent_url_final
+            final_result: Any | None = None
+            stream_idx = 0
+            async for event in transport.stream(request):
+                stream_events = stream_idx + 1
+                await _ensure_binding(
+                    context_id=event.context_id,
+                    task_id=event.task_id,
+                    agent_url_override=event.agent_url,
+                )
+                if event.agent_url is not None:
+                    remote_agent_url_final = event.agent_url
+
+                chunk_bytes = 0
+                if event.text is not None:
+                    meta = dict(event.meta) if event.meta is not None else None
+                    chunk_bytes += _text_bytes(event.text)
+                    meta_bytes = _estimate_bytes(event.meta)
+                    if meta_bytes is not None:
+                        chunk_bytes += meta_bytes
+                    await ctx.emit_chunk(
+                        parent=message,
+                        text=event.text,
+                        done=event.done,
+                        meta=meta,
+                    )
+
+                if runtime is not None:
+                    meta_keys = None
+                    if event.meta:
+                        meta_keys = sorted(event.meta.keys())
+                    extra = _merge_remote_extra(
+                        base_extra,
+                        agent_url=remote_agent_url_final,
+                        context_id=remote_context_id,
+                        task_id=remote_task_id,
+                        additional={
+                            "remote_stream_seq": stream_idx,
+                            "remote_chunk_bytes": chunk_bytes if chunk_bytes else None,
+                            "remote_chunk_done": event.done,
+                            "remote_chunk_meta_keys": meta_keys,
+                        },
+                    )
+                    await runtime.record_remote_event(
+                        event="remote_stream_event",
+                        node=owner,
+                        context=ctx,
+                        trace_id=trace_id,
+                        latency_ms=(time.perf_counter() - call_start) * 1000,
+                        level=logging.DEBUG,
+                        extra=extra,
+                    )
+
+                if chunk_bytes:
+                    response_bytes += chunk_bytes
+
+                if event.result is not None:
+                    result_bytes = _estimate_bytes(event.result)
+                    if result_bytes is not None:
+                        response_bytes += result_bytes
+                    final_result = event.result
+
+                stream_idx += 1
+
+            return final_result
+
+        call_start = time.perf_counter()
+
+        await runtime.record_remote_event(
+            event="remote_call_start",
+            node=owner,
+            context=ctx,
+            trace_id=trace_id,
+            latency_ms=0.0,
+            level=logging.DEBUG,
+            extra=_merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=None,
+                task_id=None,
+            ),
+        )
+
+        try:
+            if streaming:
+                final_result = await _run_stream()
+                result_payload = final_result
+            else:
+                result = await transport.send(request)
+                await _ensure_binding(
+                    context_id=result.context_id,
+                    task_id=result.task_id,
+                    agent_url_override=result.agent_url,
+                )
+                if result.context_id is not None:
+                    remote_context_id = result.context_id
+                if result.task_id is not None:
+                    remote_task_id = result.task_id
+                if result.agent_url is not None:
+                    remote_agent_url_final = result.agent_url
+                result_payload = result.result
+                response_size = _estimate_bytes(result_payload)
+                if response_size is not None:
+                    response_bytes += response_size
+        except TraceCancelled:
+            raise
+        except asyncio.CancelledError:
+            raise
+        except Exception as exc:
+            latency = (time.perf_counter() - call_start) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=remote_context_id,
+                task_id=remote_task_id,
+                additional={
+                    "remote_error": repr(exc),
+                    "remote_status": "error",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_error",
+                node=owner,
+                context=ctx,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.ERROR,
+                extra=extra,
+            )
+            raise
+        else:
+            latency = (time.perf_counter() - call_start) * 1000
+            extra = _merge_remote_extra(
+                base_extra,
+                agent_url=remote_agent_url_final,
+                context_id=remote_context_id,
+                task_id=remote_task_id,
+                additional={
+                    "remote_response_bytes": response_bytes,
+                    "remote_stream_events": stream_events,
+                    "remote_status": "success",
+                },
+            )
+            await runtime.record_remote_event(
+                event="remote_call_success",
+                node=owner,
+                context=ctx,
+                trace_id=trace_id,
+                latency_ms=latency,
+                level=logging.INFO,
+                extra=extra,
+            )
+            return result_payload
+        finally:
+            await _cleanup_cancel_task()
+
+    return Node(_remote_impl, name=name, policy=node_policy)
+
+
+__all__ = [
+    "RemoteCallRequest",
+    "RemoteCallResult",
+    "RemoteStreamEvent",
+    "RemoteTransport",
+    "RemoteNode",
+]

penguiflow/state.py

@@ -0,0 +1,64 @@
+"""State store protocol and helpers for PenguiFlow."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+from .metrics import FlowEvent
+
+
+@dataclass(slots=True)
+class StoredEvent:
+    """Representation of a runtime event persisted by a state store."""
+
+    trace_id: str | None
+    ts: float
+    kind: str
+    node_name: str | None
+    node_id: str | None
+    payload: Mapping[str, Any]
+
+    @classmethod
+    def from_flow_event(cls, event: FlowEvent) -> StoredEvent:
+        """Create a stored representation from a :class:`FlowEvent`."""
+
+        return cls(
+            trace_id=event.trace_id,
+            ts=event.ts,
+            kind=event.event_type,
+            node_name=event.node_name,
+            node_id=event.node_id,
+            payload=event.to_payload(),
+        )
+
+
+@dataclass(slots=True)
+class RemoteBinding:
+    """Association between a trace and a remote worker/agent."""
+
+    trace_id: str
+    context_id: str
+    task_id: str
+    agent_url: str
+
+
+class StateStore(Protocol):
+    """Protocol for durable state adapters used by PenguiFlow."""
+
+    async def save_event(self, event: StoredEvent) -> None:
+        """Persist a runtime event.
+
+        Implementations may choose any storage backend (Postgres, Redis, etc.).
+        The method must be idempotent since retries can emit duplicate events.
+        """
+
+    async def load_history(self, trace_id: str) -> Sequence[StoredEvent]:
+        """Return the ordered history for a trace id."""
+
+    async def save_remote_binding(self, binding: RemoteBinding) -> None:
+        """Persist the mapping between a trace and an external worker."""
+
+
+__all__ = ["StateStore", "StoredEvent", "RemoteBinding"]

penguiflow/testkit.py

@@ -21,9 +21,15 @@ from weakref import WeakKeyDictionary
 from .core import PenguiFlow
 from .errors import FlowErrorCode
 from .metrics import FlowEvent
-from .types import Message
+from .types import Headers, Message
 
-__all__ = ["run_one", "assert_node_sequence", "simulate_error"]
+__all__ = [
+    "run_one",
+    "assert_node_sequence",
+    "get_recorded_events",
+    "simulate_error",
+    "assert_preserves_message_envelope",
+]
 
 
 _MAX_TRACE_HISTORY = 64
@@ -102,6 +108,20 @@ class _Recorder:
         await self._state.record(event)
 
 
+class _StubContext:
+    async def emit(self, *_args: Any, **_kwargs: Any) -> None:
+        return None
+
+    def emit_nowait(self, *_args: Any, **_kwargs: Any) -> None:
+        return None
+
+    async def emit_chunk(self, *_args: Any, **_kwargs: Any) -> Any:
+        raise RuntimeError(
+            "FlowTestKit stub context does not support emit_chunk; provide a custom"
+            " context via the 'ctx' parameter"
+        )
+
+
 def _get_state(flow: PenguiFlow) -> _RecorderState:
     state = _RECORDER_STATE.get(flow)
     if state is None:
@@ -150,6 +170,78 @@ async def run_one(
     return result
 
 
+async def assert_preserves_message_envelope(
+    node: Callable[[Message, Any], Awaitable[Any]] | Any,
+    *,
+    message: Message | None = None,
+    ctx: Any | None = None,
+) -> Message:
+    """Execute ``node`` and assert it preserves the ``Message`` envelope.
+
+    Parameters
+    ----------
+    node:
+        Either a bare async callable or a :class:`penguiflow.node.Node` whose
+        first parameter is a :class:`~penguiflow.types.Message` instance.
+    message:
+        Optional sample message. When omitted, a minimal envelope is
+        synthesised.
+    ctx:
+        Optional context object passed to the node. By default a stub context
+        is used that simply no-ops ``emit``/``emit_nowait``.
+
+    Returns
+    -------
+    Message
+        The resulting message from the node, allowing additional assertions.
+
+    Raises
+    ------
+    AssertionError
+        If the node does not return a ``Message`` or mutates core envelope
+        fields (headers or trace_id).
+    TypeError
+        If ``node`` is not awaitable.
+    """
+
+    from .node import Node  # Local import to avoid circular dependency
+
+    if isinstance(node, Node):
+        func = node.func
+        node_name = node.name or node.func.__name__
+    else:
+        func = node
+        node_name = getattr(node, "__name__", "<anonymous>")
+
+    if not inspect.iscoroutinefunction(func):
+        raise TypeError("assert_preserves_message_envelope expects an async node")
+
+    sample = message or Message(payload={}, headers=Headers(tenant="test"))
+    context = ctx if ctx is not None else _StubContext()
+
+    result = await func(sample, context)
+    if not isinstance(result, Message):
+        produced = type(result).__name__
+        raise AssertionError(
+            "Node "
+            f"'{node_name}' must return a Message but produced {produced}"
+        )
+
+    mismatches: list[str] = []
+    if result.headers != sample.headers:
+        mismatches.append("headers")
+    if result.trace_id != sample.trace_id:
+        mismatches.append("trace_id")
+
+    if mismatches:
+        joined = ", ".join(mismatches)
+        raise AssertionError(
+            f"Node '{node_name}' altered Message {joined}; preserve the envelope"
+        )
+
+    return result
+
+
 def assert_node_sequence(trace_id: str, expected: Sequence[str]) -> None:
     """Assert that ``expected`` matches the recorded node start order."""
 
@@ -175,6 +267,19 @@ def assert_node_sequence(trace_id: str, expected: Sequence[str]) -> None:
         )
 
 
+def get_recorded_events(trace_id: str) -> tuple[FlowEvent, ...]:
+    """Return the recorded :class:`FlowEvent` history for ``trace_id``.
+
+    The FlowTestKit recorder maintains a bounded cache of trace histories.
+    This helper exposes the immutable snapshot so tests can assert on
+    diagnostics such as ``node_failed`` payloads or retry attempts without
+    touching the private cache directly.
+    """
+
+    events = _TRACE_HISTORY.get(trace_id, [])
+    return tuple(events)
+
+
 class _ErrorSimulation:
     def __init__(
         self,