cellaflow 0.1.0__py3-none-any.whl

cellaflow/__init__.py

@@ -0,0 +1,5 @@
+from cellaflow.client import CellaflowClient
+from cellaflow.decorators import workflow, step, tool
+from cellaflow.idempotency import IdempotencyScope
+
+__all__ = ["CellaflowClient", "workflow", "step", "tool", "IdempotencyScope"]

cellaflow/client.py

@@ -0,0 +1,165 @@
+import grpc
+from typing import Any, Dict, List, Optional, Tuple, cast
+
+from cellaflow.v1 import service_pb2, service_pb2_grpc
+from cellaflow.v1 import common_pb2, idempotency_pb2
+from cellaflow.serialization import serialize, deserialize
+
+
+class CellaflowClient:
+    """
+    gRPC Client for the Cellaflow Engine.
+    Handles communication with the engine and strictly uses
+    MessagePack for state payloads.
+    """
+
+    def __init__(self, target: str = "localhost:50051", secure: bool = False) -> None:
+        if secure:
+            self.channel = grpc.secure_channel(target, grpc.ssl_channel_credentials())
+        else:
+            self.channel = grpc.insecure_channel(target)
+        self.stub = service_pb2_grpc.WorkflowEngineServiceStub(self.channel)
+
+    def start_session(
+        self, workflow_id: str, version: str, session_id: Optional[str] = None
+    ) -> service_pb2.StartSessionResponse:
+        req = service_pb2.StartSessionRequest(
+            workflow_id=workflow_id,
+            version=version,
+        )
+        if session_id:
+            req.session_id = session_id
+
+        return cast(service_pb2.StartSessionResponse, self.stub.StartSession(req))
+
+    def commit_step(
+        self,
+        session_id: str,
+        sequence: int,
+        name: str,
+        status: "common_pb2.StepStatus.ValueType",
+        output_payload: Dict[str, Any],
+        idempotency_key: Optional[str] = None,
+        idempotency_fencing_token: Optional[int] = None,
+    ) -> service_pb2.CommitStepResponse:
+        # Strictly serialize dict to MessagePack
+        serialized_state = serialize(output_payload)
+
+        step_result = common_pb2.StepResult(
+            sequence=sequence,
+            name=name,
+            status=status,
+            output_payload=serialized_state,
+        )
+
+        req = service_pb2.CommitStepRequest(
+            session_id=session_id,
+            step_result=step_result,
+        )
+        if idempotency_key is not None:
+            req.idempotency_key = idempotency_key
+            if idempotency_fencing_token is None:
+                raise ValueError(
+                    "idempotency_fencing_token required if idempotency_key is set"
+                )
+            req.idempotency_fencing_token = idempotency_fencing_token
+
+        return cast(service_pb2.CommitStepResponse, self.stub.CommitStep(req))
+
+    def get_graph(
+        self, session_id: str, limit: Optional[int] = None, cursor: Optional[str] = None
+    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+        """
+        Returns a tuple of (list of deserialized step results, next_cursor).
+        Each step result is a dictionary representation of the StepResult proto,
+        with the output_payload fully deserialized into a Python dictionary.
+        """
+        req = service_pb2.GetGraphRequest(session_id=session_id)
+        if limit is not None:
+            req.limit = limit
+        if cursor is not None:
+            req.cursor = cursor
+
+        resp: service_pb2.GetGraphResponse = self.stub.GetGraph(req)
+
+        results = []
+        for step in resp.steps:
+            results.append(
+                {
+                    "sequence": step.sequence,
+                    "name": step.name,
+                    "status": step.status,
+                    "output_payload": deserialize(step.output_payload),
+                    "idempotency_key": (
+                        step.idempotency_key
+                        if step.HasField("idempotency_key")
+                        else None
+                    ),
+                }
+            )
+
+        next_cursor = resp.next_cursor if resp.next_cursor else None
+        return results, next_cursor
+
+    def check_idempotency_cache(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        wait_timeout_ms: Optional[int] = None,
+        lease_ttl_ms: Optional[int] = None,
+    ) -> idempotency_pb2.CheckCacheResponse:
+        req = idempotency_pb2.CheckCacheRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+        )
+        if wait_timeout_ms is not None:
+            req.wait_timeout_ms = wait_timeout_ms
+        if lease_ttl_ms is not None:
+            req.lease_ttl_ms = lease_ttl_ms
+
+        return cast(
+            idempotency_pb2.CheckCacheResponse,
+            self.stub.CheckIdempotencyCache(req),
+        )
+
+    def renew_lease(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        fencing_token: int,
+        extend_ms: int,
+    ) -> idempotency_pb2.RenewLeaseResponse:
+        req = idempotency_pb2.RenewLeaseRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+            fencing_token=fencing_token,
+            extend_ms=extend_ms,
+        )
+
+        return cast(
+            idempotency_pb2.RenewLeaseResponse,
+            self.stub.RenewLease(req),
+        )
+
+    def release_lease(
+        self,
+        agent_id: str,
+        idempotency_key: str,
+        fencing_token: int,
+        reason: Optional[str] = None,
+    ) -> idempotency_pb2.ReleaseLeaseResponse:
+        req = idempotency_pb2.ReleaseLeaseRequest(
+            agent_id=agent_id,
+            idempotency_key=idempotency_key,
+            fencing_token=fencing_token,
+        )
+        if reason:
+            req.reason = reason
+
+        return cast(
+            idempotency_pb2.ReleaseLeaseResponse,
+            self.stub.ReleaseLease(req),
+        )
+
+    def close(self) -> None:
+        self.channel.close()

cellaflow/context.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass, field
+from typing import Any, Dict
+import contextvars
+
+from cellaflow.client import CellaflowClient
+
+
+@dataclass
+class WorkflowContext:
+    client: CellaflowClient
+    session_id: str
+    workflow_version: str
+    sequence: int = 0
+    # Map sequence number to a deserialized step payload for fast replay lookups
+    replayed_steps: Dict[int, Dict[str, Any]] = field(default_factory=dict)
+
+
+_current_context: contextvars.ContextVar[WorkflowContext] = contextvars.ContextVar(
+    "workflow_context"
+)
+
+
+def get_context() -> WorkflowContext:
+    try:
+        return _current_context.get()
+    except LookupError:
+        raise RuntimeError(
+            "No active workflow context found. "
+            "Are you calling a @step inside a @workflow?"
+        )
+
+
+def set_context(context: WorkflowContext) -> contextvars.Token[WorkflowContext]:
+    return _current_context.set(context)
+
+
+def reset_context(token: contextvars.Token[WorkflowContext]) -> None:
+    _current_context.reset(token)

cellaflow/decorators.py

@@ -0,0 +1,304 @@
+import functools
+import inspect
+import uuid
+import asyncio
+import time
+import logging
+from typing import Any, Callable, TypeVar, cast, Optional
+
+from cellaflow.client import CellaflowClient
+from cellaflow.context import WorkflowContext, set_context, reset_context, get_context
+from cellaflow.idempotency import derive_idempotency_key, IdempotencyScope
+from cellaflow.lease import LeaseHeartbeat
+from cellaflow.serialization import deserialize
+from cellaflow.v1.common_pb2 import STEP_STATUS_SUCCESS
+from cellaflow.v1.idempotency_pb2 import (
+    CACHE_STATUS_HIT,
+    CACHE_STATUS_ACQUIRED,
+    CACHE_STATUS_IN_PROGRESS,
+)
+
+logger = logging.getLogger(__name__)
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def workflow(
+    version: str = "1.0.0", target: str = "localhost:50051", secure: bool = False
+) -> Callable[[F], F]:
+    def decorator(func: F) -> F:
+        @functools.wraps(func)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            # Pop _session_id if provided for recovery, else generate new
+            session_id = kwargs.pop("_session_id", str(uuid.uuid4()))
+
+            client = CellaflowClient(target=target, secure=secure)
+            resp = client.start_session(
+                workflow_id=func.__name__, version=version, session_id=session_id
+            )
+
+            replayed_steps = {}
+            if resp.is_recovered:
+                steps, _ = client.get_graph(resp.session_id)
+                for step_info in steps:
+                    replayed_steps[step_info["sequence"]] = step_info["output_payload"]
+
+            ctx = WorkflowContext(
+                client=client,
+                session_id=resp.session_id,
+                workflow_version=resp.version,
+                sequence=0,
+                replayed_steps=replayed_steps,
+            )
+            token = set_context(ctx)
+            try:
+                return await func(*args, **kwargs)
+            finally:
+                reset_context(token)
+
+        @functools.wraps(func)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            session_id = kwargs.pop("_session_id", str(uuid.uuid4()))
+
+            client = CellaflowClient(target=target, secure=secure)
+            resp = client.start_session(
+                workflow_id=func.__name__, version=version, session_id=session_id
+            )
+
+            replayed_steps = {}
+            if resp.is_recovered:
+                steps, _ = client.get_graph(resp.session_id)
+                for step_info in steps:
+                    replayed_steps[step_info["sequence"]] = step_info["output_payload"]
+
+            ctx = WorkflowContext(
+                client=client,
+                session_id=resp.session_id,
+                workflow_version=resp.version,
+                sequence=0,
+                replayed_steps=replayed_steps,
+            )
+            token = set_context(ctx)
+            try:
+                return func(*args, **kwargs)
+            finally:
+                reset_context(token)
+
+        if inspect.iscoroutinefunction(func):
+            return cast(F, async_wrapper)
+        return cast(F, sync_wrapper)
+
+    return decorator
+
+
+def step(
+    func: Optional[F] = None,
+    *,
+    idempotency_key: Optional[str] = None,
+    agent_id: str = "default",
+    tool_name: Optional[str] = None,
+    scope: IdempotencyScope = IdempotencyScope.SCOPE_SESSION_WIDE,
+) -> Any:
+    if func is None:
+        return functools.partial(
+            step,
+            idempotency_key=idempotency_key,
+            agent_id=agent_id,
+            tool_name=tool_name,
+            scope=scope,
+        )
+
+    actual_tool_name = tool_name or func.__name__
+
+    @functools.wraps(func)
+    async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+        ctx = get_context()
+        ctx.sequence += 1
+        seq = ctx.sequence
+
+        if seq in ctx.replayed_steps:
+            return ctx.replayed_steps[seq].get("result")
+
+        # Derive idempotency key
+        ikey = idempotency_key
+        if not ikey:
+            ikey = derive_idempotency_key(
+                ctx.session_id,
+                ctx.workflow_version,
+                seq,
+                agent_id,
+                actual_tool_name,
+                scope,
+                *args,
+                **kwargs,
+            )
+
+        fencing_token = 0
+        hb: Optional[LeaseHeartbeat] = None
+
+        while True:
+            resp = ctx.client.check_idempotency_cache(
+                agent_id=agent_id,
+                idempotency_key=ikey,
+            )
+            if resp.status == CACHE_STATUS_HIT:
+                if resp.cached_result and resp.cached_result.output_payload:
+                    deserialized = deserialize(resp.cached_result.output_payload)
+                    return deserialized.get("result")
+                return None
+            elif resp.status == CACHE_STATUS_IN_PROGRESS:
+                retry_ms = resp.retry_after_ms or 1000
+                logger.info(
+                    "Step %s is in progress by another worker. Sleeping %d ms...",
+                    actual_tool_name,
+                    retry_ms,
+                )
+                await asyncio.sleep(retry_ms / 1000.0)
+            elif resp.status == CACHE_STATUS_ACQUIRED:
+                fencing_token = resp.fencing_token or 0
+                interval_ms = resp.heartbeat_interval_ms or 5000
+                hb = LeaseHeartbeat(
+                    client=ctx.client,
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    heartbeat_interval_ms=interval_ms,
+                )
+                hb.start_async()
+                break
+            else:
+                break
+
+        try:
+            result = await func(*args, **kwargs)
+            ctx.client.commit_step(
+                session_id=ctx.session_id,
+                sequence=seq,
+                name=actual_tool_name,
+                status=STEP_STATUS_SUCCESS,
+                output_payload={"result": result},
+                idempotency_key=ikey,
+                idempotency_fencing_token=fencing_token,
+            )
+            return result
+        except Exception as e:
+            if fencing_token > 0:
+                ctx.client.release_lease(
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    reason="TOOL_ERROR",
+                )
+            raise e
+        finally:
+            if hb:
+                await hb.stop_async()
+
+    @functools.wraps(func)
+    def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+        ctx = get_context()
+        ctx.sequence += 1
+        seq = ctx.sequence
+
+        if seq in ctx.replayed_steps:
+            return ctx.replayed_steps[seq].get("result")
+
+        # Derive idempotency key
+        ikey = idempotency_key
+        if not ikey:
+            ikey = derive_idempotency_key(
+                ctx.session_id,
+                ctx.workflow_version,
+                seq,
+                agent_id,
+                actual_tool_name,
+                scope,
+                *args,
+                **kwargs,
+            )
+
+        fencing_token = 0
+        hb: Optional[LeaseHeartbeat] = None
+
+        while True:
+            resp = ctx.client.check_idempotency_cache(
+                agent_id=agent_id,
+                idempotency_key=ikey,
+            )
+            if resp.status == CACHE_STATUS_HIT:
+                if resp.cached_result and resp.cached_result.output_payload:
+                    deserialized = deserialize(resp.cached_result.output_payload)
+                    return deserialized.get("result")
+                return None
+            elif resp.status == CACHE_STATUS_IN_PROGRESS:
+                retry_ms = resp.retry_after_ms or 1000
+                logger.info(
+                    "Step %s is in progress by another worker. Sleeping %d ms...",
+                    actual_tool_name,
+                    retry_ms,
+                )
+                time.sleep(retry_ms / 1000.0)
+            elif resp.status == CACHE_STATUS_ACQUIRED:
+                fencing_token = resp.fencing_token or 0
+                interval_ms = resp.heartbeat_interval_ms or 5000
+                hb = LeaseHeartbeat(
+                    client=ctx.client,
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    heartbeat_interval_ms=interval_ms,
+                )
+                hb.start_sync()
+                break
+            else:
+                break
+
+        try:
+            result = func(*args, **kwargs)
+            ctx.client.commit_step(
+                session_id=ctx.session_id,
+                sequence=seq,
+                name=actual_tool_name,
+                status=STEP_STATUS_SUCCESS,
+                output_payload={"result": result},
+                idempotency_key=ikey,
+                idempotency_fencing_token=fencing_token,
+            )
+            return result
+        except Exception as e:
+            if fencing_token > 0:
+                ctx.client.release_lease(
+                    agent_id=agent_id,
+                    idempotency_key=ikey,
+                    fencing_token=fencing_token,
+                    reason="TOOL_ERROR",
+                )
+            raise e
+        finally:
+            if hb:
+                hb.stop_sync()
+
+    if inspect.iscoroutinefunction(func):
+        return cast(F, async_wrapper)
+    return cast(F, sync_wrapper)
+
+
+def tool(
+    func: Optional[F] = None,
+    *,
+    idempotency_key: Optional[str] = None,
+    agent_id: str = "default",
+    tool_name: Optional[str] = None,
+    scope: IdempotencyScope = IdempotencyScope.SCOPE_SESSION_WIDE,
+) -> Any:
+    """
+    Decorator for tool steps that require idempotency tracking.
+    Functionally identical to @step for now, but semantically distinct.
+    """
+    return step(
+        func=func,
+        idempotency_key=idempotency_key,
+        agent_id=agent_id,
+        tool_name=tool_name,
+        scope=scope,
+    )

cellaflow/idempotency.py

@@ -0,0 +1,61 @@
+import hashlib
+import rfc8785
+from typing import Any
+
+
+from enum import Enum
+
+
+class IdempotencyScope(Enum):
+    SCOPE_UNSPECIFIED = 0
+    SCOPE_SESSION_WIDE = 1  # Shared across all agents in the session (Default)
+    SCOPE_AGENT_PRIVATE = 2  # Isolated to the executing agent
+    SCOPE_STEP_LOCAL = 3  # Isolated to the specific superstep / node
+
+
+def _hash_inputs(*args: Any, **kwargs: Any) -> str:
+    """
+    Hashes the inputs using RFC 8785 Canonical JSON and SHA-256.
+    Returns the first 16 bytes of the hash, hex-encoded (32 characters).
+    """
+    payload = {"args": list(args), "kwargs": kwargs}
+
+    # Serialize to canonical JSON bytes according to RFC 8785
+    canonical_bytes = rfc8785.dumps(payload)  # type: ignore[arg-type]
+
+    # Hash using SHA-256
+    hash_obj = hashlib.sha256(canonical_bytes)
+
+    # Return the first 16 bytes hex-encoded
+    return hash_obj.digest()[:16].hex()
+
+
+def derive_idempotency_key(
+    session_id: str,
+    workflow_version: str,
+    step_sequence: int,
+    agent_id: str,
+    tool_name: str,
+    scope: IdempotencyScope,
+    *args: Any,
+    **kwargs: Any,
+) -> str:
+    """
+    Derives the canonical idempotency key for a step or tool execution.
+    Format varies based on IdempotencyScope.
+    """
+    inputs_hash = _hash_inputs(*args, **kwargs)
+
+    seq_part = "session_wide"
+    agent_part = "session_wide"
+
+    if scope == IdempotencyScope.SCOPE_AGENT_PRIVATE:
+        agent_part = agent_id
+    elif scope == IdempotencyScope.SCOPE_STEP_LOCAL:
+        seq_part = str(step_sequence)
+        agent_part = agent_id
+
+    return (
+        f"{session_id}:{workflow_version}:{seq_part}:"
+        f"{agent_part}:{tool_name}:{inputs_hash}"
+    )

cellaflow/lease.py

@@ -0,0 +1,120 @@
+import asyncio
+import threading
+import logging
+from typing import Optional, Any
+from cellaflow.client import CellaflowClient
+
+logger = logging.getLogger(__name__)
+
+
+class LeaseHeartbeat:
+    """
+    Manages the background heartbeat (RenewLease) for an actively running step.
+    Supports both sync (daemon thread) and async (asyncio task) environments.
+    """
+
+    def __init__(
+        self,
+        client: CellaflowClient,
+        agent_id: str,
+        idempotency_key: str,
+        fencing_token: int,
+        heartbeat_interval_ms: int,
+    ) -> None:
+        self.client = client
+        self.agent_id = agent_id
+        self.idempotency_key = idempotency_key
+        self.fencing_token = fencing_token
+        # Renew slightly before expiration
+        self.interval_sec = max(0.1, (heartbeat_interval_ms - 100) / 1000.0)
+        self.extend_ms = (
+            heartbeat_interval_ms * 4
+        )  # Usually TTL is 4x heartbeat interval
+
+        # State tracking
+        self._stop_event_sync = threading.Event()
+        self._sync_thread: Optional[threading.Thread] = None
+
+        self._stop_event_async: Optional[asyncio.Event] = None
+        self._async_task: Optional[asyncio.Task[Any]] = None
+
+    def start_sync(self) -> None:
+        """Starts a daemon thread for synchronous execution."""
+        self._stop_event_sync.clear()
+        self._sync_thread = threading.Thread(
+            target=self._sync_loop, daemon=True, name=f"lease-hb-{self.fencing_token}"
+        )
+        self._sync_thread.start()
+
+    def stop_sync(self) -> None:
+        """Stops the daemon thread cleanly."""
+        self._stop_event_sync.set()
+        if self._sync_thread and self._sync_thread.is_alive():
+            self._sync_thread.join(timeout=1.0)
+
+    def _sync_loop(self) -> None:
+        while not self._stop_event_sync.wait(timeout=self.interval_sec):
+            try:
+                resp = self.client.renew_lease(
+                    agent_id=self.agent_id,
+                    idempotency_key=self.idempotency_key,
+                    fencing_token=self.fencing_token,
+                    extend_ms=self.extend_ms,
+                )
+                if not resp.renewed:
+                    logger.warning(
+                        f"Lease {self.idempotency_key} failed to renew: "
+                        f"{resp.failure_reason}"
+                    )
+                    break
+            except Exception as e:
+                logger.error(f"Lease renewal error: {e}")
+
+    def start_async(self) -> None:
+        """Starts an asyncio task for asynchronous execution."""
+        self._stop_event_async = asyncio.Event()
+        loop = asyncio.get_running_loop()
+        self._async_task = loop.create_task(self._async_loop())
+
+    async def stop_async(self) -> None:
+        """Stops the asyncio task cleanly."""
+        if self._stop_event_async:
+            self._stop_event_async.set()
+        if self._async_task and not self._async_task.done():
+            self._async_task.cancel()
+            try:
+                await self._async_task
+            except asyncio.CancelledError:
+                pass
+
+    async def _async_loop(self) -> None:
+        if not self._stop_event_async:
+            return
+
+        # We use asyncio.wait_for to wait for the stop event, with a timeout
+        while True:
+            try:
+                await asyncio.wait_for(
+                    self._stop_event_async.wait(), timeout=self.interval_sec
+                )
+                # If we get here, the stop event was set
+                break
+            except asyncio.TimeoutError:
+                # Timeout means interval elapsed, do heartbeat
+                try:
+                    resp = self.client.renew_lease(
+                        agent_id=self.agent_id,
+                        idempotency_key=self.idempotency_key,
+                        fencing_token=self.fencing_token,
+                        extend_ms=self.extend_ms,
+                    )
+                    if not resp.renewed:
+                        logger.warning(
+                            f"Lease {self.idempotency_key} failed to renew: "
+                            f"{resp.failure_reason}"
+                        )
+                        break
+                except Exception as e:
+                    logger.error(f"Lease renewal error: {e}")
+            except asyncio.CancelledError:
+                break