by-framework 0.2.2.dev0__py3-none-any.whl → 0.2.2.dev2__py3-none-any.whl

by_framework/__init__.py

@@ -8,6 +8,7 @@ from `GatewayWorker` and running `run_worker`.
 from .client.byai_client import ByaiGatewayClient
 from .client.client import (
     CancelTaskResponse,
+    DataStreamEntry,
     GatewayClient,
     GatewayInterceptor,
     SendMessageResponse,
@@ -140,6 +141,7 @@ __all__ = [
     "GatewayClient",
     "ByaiGatewayClient",
     "GatewayInterceptor",
+    "DataStreamEntry",
     "SendMessageResponse",
     "CancelTaskResponse",
     "run_worker",

by_framework/client/__init__.py

@@ -1,10 +1,11 @@
 """Client module for Gateway communication."""
 
 from .byai_client import ByaiGatewayClient
-from .client import GatewayClient, GatewayInterceptor
+from .client import DataStreamEntry, GatewayClient, GatewayInterceptor
 
 __all__ = [
     "GatewayClient",
     "ByaiGatewayClient",
     "GatewayInterceptor",
+    "DataStreamEntry",
 ]

by_framework/client/client.py

@@ -9,7 +9,8 @@ import json
 import time
 import uuid
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Protocol
+from dataclasses import fields as dataclass_fields
+from typing import (TYPE_CHECKING, Any, AsyncIterator, Dict, List, Optional, Protocol)
 
 from by_framework.common.constants import (
     CANCEL_MESSAGE_ID_PREFIX,
@@ -17,6 +18,7 @@ from by_framework.common.constants import (
     MESSAGE_ID_PREFIX,
     RedisKeys,
 )
+from by_framework.common.logger import logger
 from by_framework.common.redis_client import Redis, get_redis
 from by_framework.core.availability import (
     AvailabilityRouter,
@@ -32,6 +34,7 @@ from by_framework.core.protocol.commands import (
     ReloadPluginsCommand,
     ResumeCommand,
 )
+from by_framework.core.protocol.data_message import DataMessage
 from by_framework.core.protocol.message_header import MessageHeader
 from by_framework.core.protocol.responses import (
     CancelTaskResponse,
@@ -40,6 +43,7 @@ from by_framework.core.protocol.responses import (
 )
 from by_framework.core.registry import WorkerRegistry
 from by_framework.errors import WorkerRegistryNotSetError
+from by_framework.trace.span_recorder import (SpanRecorder, TraceSpan, str_to_uint64)
 
 if TYPE_CHECKING:
     pass
@@ -59,6 +63,14 @@ class RouteResolution:
     target_worker_id: str = ""
 
 
+@dataclass(frozen=True)
+class DataStreamEntry:
+    """A decoded entry from a session data stream."""
+
+    stream_id: str
+    message: DataMessage
+
+
 class GatewayClient:
     """Gateway client for sending messages and cancel requests to Gateway workers.
 
@@ -76,16 +88,183 @@ class GatewayClient:
         registry: Optional[WorkerRegistry] = None,
         redis_client: Optional[Redis] = None,
         interceptors: Optional[List[GatewayInterceptor]] = None,
+        span_recorder: Optional[SpanRecorder] = None,
     ):
         self.registry = registry
         self.redis = (
             redis_client or (registry.redis if registry else None) or get_redis()
         )
         self.interceptors = interceptors or []
+        self.span_recorder = span_recorder or SpanRecorder(self.redis)
 
     def add_interceptor(self, interceptor: GatewayInterceptor):
         self.interceptors.append(interceptor)
 
+    @staticmethod
+    def _decode_redis_value(value: Any) -> Any:
+        """Decode Redis bytes values while preserving already-decoded clients."""
+        if isinstance(value, bytes):
+            return value.decode("utf-8")
+        return value
+
+    @classmethod
+    def _decode_data_stream_entry(
+        cls, stream_id: Any, fields: Dict[Any, Any]
+    ) -> DataStreamEntry:
+        raw = fields.get(b"data")
+        if raw is None:
+            raw = fields.get("data")
+        if raw is None:
+            raise ValueError("data stream entry missing 'data' field")
+
+        payload = json.loads(cls._decode_redis_value(raw))
+        data_message_fields = {field.name for field in dataclass_fields(DataMessage)}
+        return DataStreamEntry(
+            stream_id=cls._decode_redis_value(stream_id),
+            message=DataMessage(
+                **{
+                    key: value
+                    for key, value in payload.items()
+                    if key in data_message_fields
+                }
+            ),
+        )
+
+    async def read_data_messages(
+        self,
+        session_id: str,
+        last_id: str = "0-0",
+        block_ms: int = 0,
+        count: int = 100,
+    ) -> List[DataStreamEntry]:
+        """Read decoded messages from the session data stream.
+
+        Pass the last returned ``stream_id`` as ``last_id`` to continue from
+        the next entry. ``block_ms`` is passed to Redis XREAD; ``0`` means
+        block indefinitely on standard Redis clients.
+        """
+        stream_name = RedisKeys.session_data_stream(session_id)
+        messages = await self.redis.xread(
+            streams={stream_name: last_id},
+            count=count,
+            block=block_ms,
+        )
+
+        results: List[DataStreamEntry] = []
+        for _, msg_list in messages or []:
+            for stream_id, fields in msg_list:
+                results.append(self._decode_data_stream_entry(stream_id, fields))
+        return results
+
+    async def get_data_message_checkpoint(
+        self,
+        session_id: str,
+        consumer_name: str,
+    ) -> str:
+        """Return the last committed data stream ID for a named consumer."""
+        checkpoint = await self.redis.get(
+            RedisKeys.session_data_checkpoint(session_id, consumer_name)
+        )
+        if checkpoint is None:
+            return "0-0"
+        return self._decode_redis_value(checkpoint)
+
+    async def commit_data_message(
+        self,
+        session_id: str,
+        stream_id: str,
+        consumer_name: str,
+    ) -> None:
+        """Commit a data stream ID as processed for a named consumer."""
+        await self.redis.set(
+            RedisKeys.session_data_checkpoint(session_id, consumer_name),
+            stream_id,
+            ex=RedisKeys.DEFAULT_SESSION_TTL,
+        )
+
+    async def read_data_messages_from_checkpoint(
+        self,
+        session_id: str,
+        consumer_name: str,
+        block_ms: int = 0,
+        count: int = 100,
+        auto_commit: bool = False,
+    ) -> List[DataStreamEntry]:
+        """Read messages starting after a named consumer's committed checkpoint."""
+        last_id = await self.get_data_message_checkpoint(session_id, consumer_name)
+        entries = await self.read_data_messages(
+            session_id=session_id,
+            last_id=last_id,
+            block_ms=block_ms,
+            count=count,
+        )
+        if auto_commit and entries:
+            await self.commit_data_message(
+                session_id=session_id,
+                stream_id=entries[-1].stream_id,
+                consumer_name=consumer_name,
+            )
+        return entries
+
+    async def iter_data_messages(
+        self,
+        session_id: str,
+        last_id: str = "$",
+        block_ms: int = 5000,
+        count: int = 100,
+    ) -> AsyncIterator[DataStreamEntry]:
+        """Continuously consume decoded messages from the session data stream.
+
+        The iterator does not stop on its own. Callers should break when their
+        business-level terminal event is observed.
+        """
+        current_id = last_id
+        while True:
+            entries = await self.read_data_messages(
+                session_id=session_id,
+                last_id=current_id,
+                block_ms=block_ms,
+                count=count,
+            )
+            for entry in entries:
+                current_id = entry.stream_id
+                yield entry
+
+    async def consume_data_messages(
+        self,
+        session_id: str,
+        consumer_name: str,
+        block_ms: int = 5000,
+        count: int = 100,
+    ) -> AsyncIterator[DataStreamEntry]:
+        """Continuously consume data stream messages with checkpoint commits.
+
+        Each entry is committed after the caller's loop body completes and asks
+        for the next item. If processing fails or the iterator is closed before
+        the next item, the current entry is not committed and will be retried
+        from the checkpoint on the next consumer run. The iterator does not stop
+        on its own; callers should break on their terminal event.
+        """
+        current_id = await self.get_data_message_checkpoint(
+            session_id=session_id,
+            consumer_name=consumer_name,
+        )
+        while True:
+            entries = await self.read_data_messages(
+                session_id=session_id,
+                last_id=current_id,
+                block_ms=block_ms,
+                count=count,
+            )
+            for entry in entries:
+                yield entry
+                await self.commit_data_message(
+                    session_id=session_id,
+                    stream_id=entry.stream_id,
+                    consumer_name=consumer_name,
+                )
+                current_id = entry.stream_id
+
     async def reload_plugins_for_agent_type(
         self,
         agent_type: str,
@@ -341,11 +520,16 @@ class GatewayClient:
             )
 
             if node_worker_id:
+                node_trace_id = (
+                    node.get("trace_id")
+                    or execution.get("trace_id")
+                    or uuid.uuid4().hex
+                )
                 cancel_command = CancelTaskCommand(
                     header=MessageHeader(
                         message_id=f"{CANCEL_MESSAGE_ID_PREFIX}{uuid.uuid4().hex[:8]}",
                         session_id=session_id,
-                        trace_id=uuid.uuid4().hex,
+                        trace_id=node_trace_id,
                         target_agent_type=node.get("target_agent_type", ""),
                         parent_message_id=node_message_id,
                     ),
@@ -424,6 +608,32 @@ class GatewayClient:
         if not trace_id:
             trace_id = uuid.uuid4().hex
 
+        metadata = dict(params.get("metadata", {}) or {})
+        trace_parent_span_id = metadata.pop("trace_parent_span_id", "")
+        langfuse_parent_observation_id = metadata.pop(
+            "langfuse_parent_observation_id", ""
+        )
+        if not trace_parent_span_id:
+            trace_parent_span_id = (
+                f"{str_to_uint64(f'{message_id}:client.dispatch'):016x}"
+            )
+
+        langfuse_client_dispatch = None
+        if not params["parent_message_id"]:
+            langfuse_client_dispatch = self._start_langfuse_client_dispatch_observation(
+                trace_id=trace_id,
+                message_id=message_id,
+                target_agent_type=params["target_agent_type"],
+                session_id=params["session_id"],
+                user_code=params["user_code"],
+                user_name=params["user_name"],
+                content=params["content"],
+                metadata=metadata,
+            )
+            observation_id = getattr(langfuse_client_dispatch, "id", "")
+            if observation_id:
+                langfuse_parent_observation_id = observation_id
+
         header = MessageHeader(
             message_id=message_id,
             session_id=params["session_id"],
@@ -432,7 +642,9 @@ class GatewayClient:
             parent_message_id=params["parent_message_id"],
             user_code=params["user_code"],
             user_name=params["user_name"],
-            metadata=params["metadata"],
+            metadata=metadata,
+            trace_parent_span_id=trace_parent_span_id,
+            langfuse_parent_observation_id=langfuse_parent_observation_id,
         )
         command = self._build_gateway_command(
             action_type=params["action_type"],
@@ -451,6 +663,7 @@ class GatewayClient:
                     route_policy != RoutePolicy.SEND_ANYWAY,
                 )
             else:
+                avail_start_ms = int(time.time() * 1000)
                 availability = await AvailabilityRouter(
                     self.redis, self.registry
                 ).prepare_delivery(
@@ -470,13 +683,42 @@ class GatewayClient:
                         metadata=params["metadata"],
                     )
                 )
+                try:
+                    from by_framework.metrics import record_availability_metrics
+
+                    record_availability_metrics(
+                        agent_type=params["target_agent_type"],
+                        policy=route_policy,
+                        status=availability.status,
+                        routing_ms=float(int(time.time() * 1000) - avail_start_ms),
+                    )
+                except Exception:  # pylint: disable=broad-exception-caught
+                    pass
                 if availability.status not in (
                     AvailabilityStatus.DELIVER_NOW,
                     AvailabilityStatus.WAIT_AND_DELIVER,
                     AvailabilityStatus.FALLBACK_TO_OTHER_AGENT_TYPE,
                     AvailabilityStatus.QUEUE_PENDING,
                 ):
-                    return SendMessageResponse(
+                    if self.registry and hasattr(
+                        self.registry, "record_failed_route_decision"
+                    ):
+                        await self.registry.record_failed_route_decision(
+                            execution_id=execution_id,
+                            message_id=message_id,
+                            session_id=params["session_id"],
+                            trace_id=trace_id,
+                            target_agent_type=params["target_agent_type"],
+                            parent_message_id=params["parent_message_id"] or "",
+                            source_agent_type="client",
+                            route_policy=route_policy,
+                            route_status=availability.status,
+                            stream_name=availability.stream_name or "",
+                            selected_agent_type=availability.selected_agent_type or "",
+                            availability_error_code=availability.error_code or "",
+                            availability_error=availability.error or "",
+                        )
+                    response = SendMessageResponse(
                         success=False,
                         status=ExecutionStatus.FAILED,
                         message_id="",
@@ -487,6 +729,12 @@ class GatewayClient:
                         error_code=availability.error_code
                         or ExecutionStatus.ERR_AGENT_TYPE_UNAVAILABLE,
                     )
+                    self._end_langfuse_client_dispatch_observation(
+                        langfuse_client_dispatch,
+                        output={"success": False, "error": availability.error},
+                        error=availability.error,
+                    )
+                    return response
                 if availability.status == AvailabilityStatus.QUEUE_PENDING:
                     should_dispatch_control = False
                 route = RouteResolution(
@@ -497,6 +745,11 @@ class GatewayClient:
                     params["target_agent_type"] = availability.selected_agent_type
                     command.header.target_agent_type = availability.selected_agent_type
         except LookupError as err:
+            self._end_langfuse_client_dispatch_observation(
+                langfuse_client_dispatch,
+                output={"success": False, "error": str(err)},
+                error=str(err),
+            )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -508,6 +761,11 @@ class GatewayClient:
                 error_code=ExecutionStatus.ERR_WORKER_NOT_ONLINE,
             )
         except ValueError as err:
+            self._end_langfuse_client_dispatch_observation(
+                langfuse_client_dispatch,
+                output={"success": False, "error": str(err)},
+                error=str(err),
+            )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -533,16 +791,44 @@ class GatewayClient:
                         "target_agent_type": params["target_agent_type"],
                         "stream_name": route.stream_name,
                         "status": "QUEUED",
+                        "route_policy": route_policy,
+                        "route_status": availability.status
+                        if not target_worker_id
+                        else "DIRECT_WORKER",
+                        "selected_agent_type": availability.selected_agent_type
+                        if not target_worker_id
+                        else "",
+                        "availability_error_code": availability.error_code
+                        if not target_worker_id
+                        else "",
+                        "availability_error": availability.error
+                        if not target_worker_id
+                        else "",
                     }
                 )
             except Exception:  # pylint: disable=broad-exception-caught
                 pass  # Fallback if registry fails
 
         # 4. Route to the appropriate stream
+        dispatch_started_at = int(time.time() * 1000)
         if should_dispatch_control:
             await self.redis.xadd(route.stream_name, command.to_redis_payload())
+        await self._record_client_dispatch_span(
+            trace_id=trace_id,
+            message_id=message_id,
+            session_id=params["session_id"],
+            parent_message_id=params["parent_message_id"] or "",
+            target_agent_type=params["target_agent_type"],
+            target_worker_id=route.target_worker_id,
+            route_policy=route_policy,
+            route_status=availability.status
+            if not target_worker_id
+            else "DIRECT_WORKER",
+            start_ts=dispatch_started_at,
+            end_ts=int(time.time() * 1000),
+        )
 
-        return SendMessageResponse(
+        response = SendMessageResponse(
             success=True,
             message_id=message_id,
             trace_id=trace_id,
@@ -550,3 +836,118 @@ class GatewayClient:
             timestamp=int(time.time() * 1000),
             status=ExecutionStatus.QUEUED,
         )
+        self._end_langfuse_client_dispatch_observation(
+            langfuse_client_dispatch,
+            output={
+                "success": True,
+                "message_id": message_id,
+                "trace_id": trace_id,
+                "target_worker_id": route.target_worker_id,
+                "status": response.status,
+            },
+        )
+        return response
+
+    def _start_langfuse_client_dispatch_observation(
+        self,
+        *,
+        trace_id: str,
+        message_id: str,
+        target_agent_type: str,
+        session_id: str,
+        user_code: str,
+        user_name: str,
+        content: Any,
+        metadata: Dict[str, Any],
+    ) -> Any:
+        try:
+            from by_framework_trace_langfuse import start_client_dispatch_observation
+
+            return start_client_dispatch_observation(
+                trace_id=trace_id,
+                message_id=message_id,
+                target_agent_type=target_agent_type,
+                session_id=session_id,
+                user_code=user_code,
+                user_name=user_name,
+                content=content,
+                metadata=metadata,
+            )
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            logger.warning(
+                "Langfuse client.dispatch observation skipped: %s",
+                err,
+                exc_info=True,
+            )
+            return None
+
+    @staticmethod
+    def _end_langfuse_client_dispatch_observation(
+        observation: Any,
+        *,
+        output: Any,
+        error: str = "",
+    ) -> None:
+        if observation is None:
+            return
+        try:
+            if error and hasattr(observation, "update"):
+                observation.update(level="ERROR", status_message=error)
+            observation.end(output=output)
+        except TypeError:
+            try:
+                observation.update(output=output)
+                observation.end()
+            except Exception:  # pylint: disable=broad-exception-caught
+                pass
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    async def _record_client_dispatch_span(
+        self,
+        *,
+        trace_id: str,
+        message_id: str,
+        session_id: str,
+        parent_message_id: str,
+        target_agent_type: str,
+        target_worker_id: str,
+        route_policy: str,
+        route_status: str,
+        start_ts: int,
+        end_ts: int,
+    ) -> None:
+        try:
+            logger.info(
+                "Recording client dispatch span: message_id=%s, trace_id=%s",
+                message_id,
+                trace_id,
+            )
+            await self.span_recorder.record_span(
+                TraceSpan(
+                    trace_id=trace_id,
+                    span_id=f"{message_id}:client.dispatch",
+                    parent_span_id="",
+                    operation="client.dispatch",
+                    component="client",
+                    start_ts=start_ts,
+                    end_ts=end_ts,
+                    status="COMPLETED",
+                    session_id=session_id,
+                    message_id=message_id,
+                    parent_message_id=parent_message_id,
+                    worker_id=target_worker_id,
+                    source_agent_type="client",
+                    target_agent_type=target_agent_type,
+                    route_policy=route_policy,
+                    route_status=route_status,
+                )
+            )
+            logger.info(
+                "Client dispatch span recorded successfully for message_id=%s",
+                message_id,
+            )
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            logger.warning(
+                "Failed to record client dispatch span: %s", err, exc_info=True
+            )

by_framework/common/constants.py

@@ -92,6 +92,36 @@ class RedisKeys:
         """Session-level data stream. Workers push streaming content here."""
         return f"byai_gateway:session:{session_id}:data_stream"
 
+    @staticmethod
+    def session_data_checkpoint(session_id: str, consumer_name: str) -> str:
+        """Checkpoint key storing a consumer's last processed data stream ID."""
+        return f"byai_gateway:session:{session_id}:consumer:{consumer_name}:checkpoint"
+
+    @staticmethod
+    def trace_meta(trace_id: str) -> str:
+        """Hash storing trace-level metadata for observability."""
+        return f"by_framework:trace:{trace_id}"
+
+    @staticmethod
+    def trace_spans(trace_id: str) -> str:
+        """List storing trace span JSON payloads ordered by write time."""
+        return f"by_framework:trace:spans:{trace_id}"
+
+    @staticmethod
+    def trace_index_session(session_id: str) -> str:
+        """Sorted Set index from session_id to trace IDs."""
+        return f"by_framework:trace:idx:session:{session_id}"
+
+    @staticmethod
+    def trace_index_worker(worker_id: str) -> str:
+        """Sorted Set index from worker_id to trace IDs."""
+        return f"by_framework:trace:idx:worker:{worker_id}"
+
+    @staticmethod
+    def trace_index_agent(agent_type: str) -> str:
+        """Sorted Set index from agent type to trace IDs."""
+        return f"by_framework:trace:idx:agent:{agent_type}"
+
     @staticmethod
     def task_group(group_id: str) -> str:
         """Task group progress tracking Hash Key."""