by-framework 0.2.2.dev3__py3-none-any.whl → 0.2.2.dev5__py3-none-any.whl

by_framework/client/client.py

@@ -43,7 +43,16 @@ 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)
+from by_framework.trace.span_recorder import (
+    ObservabilityConfig,
+    SpanRecorder,
+    TraceSpan,
+    build_observability_config,
+    sanitize_io_value,
+    str_to_uint64,
+)
+from by_framework.trace.trace_schema import TraceRecord
+from by_framework.trace.trace_writer import TraceWriteClient
 
 if TYPE_CHECKING:
     pass
@@ -89,13 +98,30 @@ class GatewayClient:
         redis_client: Optional[Redis] = None,
         interceptors: Optional[List[GatewayInterceptor]] = None,
         span_recorder: Optional[SpanRecorder] = None,
+        obs_config: Optional[ObservabilityConfig] = 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)
+        self._obs_config = obs_config or build_observability_config()
+        self.span_recorder = span_recorder or SpanRecorder(
+            self.redis, config=self._obs_config
+        )
+        self._trace_writer = TraceWriteClient(
+            self.redis, ttl_seconds=self._obs_config.ttl_seconds
+        )
+        self._langfuse_dispatch_fn = self._resolve_langfuse_dispatch_fn()
+
+    @staticmethod
+    def _resolve_langfuse_dispatch_fn() -> Any:
+        try:
+            from by_framework_trace_langfuse import start_client_dispatch_observation
+
+            return start_client_dispatch_observation
+        except ImportError:
+            return None
 
     def add_interceptor(self, interceptor: GatewayInterceptor):
         self.interceptors.append(interceptor)
@@ -607,6 +633,20 @@ class GatewayClient:
             message_id = f"{MESSAGE_ID_PREFIX}{uuid.uuid4().hex[:8]}"
         if not trace_id:
             trace_id = uuid.uuid4().hex
+        trace_start_ts = int(time.time() * 1000)
+        is_root_dispatch = not params.get("parent_message_id", "")
+
+        # Write trace root start so the trace is visible even before the worker
+        # picks up the task.  Only written for root dispatches (no parent).
+        if is_root_dispatch:
+            await self._write_trace_root_start(
+                trace_id=trace_id,
+                message_id=message_id,
+                session_id=str(params.get("session_id", "")),
+                target_agent_type=str(params.get("target_agent_type", "")),
+                content=params.get("content"),
+                start_ts=trace_start_ts,
+            )
 
         metadata = dict(params.get("metadata", {}) or {})
         trace_parent_span_id = metadata.pop("trace_parent_span_id", "")
@@ -619,7 +659,7 @@ class GatewayClient:
             )
 
         langfuse_client_dispatch = None
-        if not params["parent_message_id"]:
+        if not langfuse_parent_observation_id:
             langfuse_client_dispatch = self._start_langfuse_client_dispatch_observation(
                 trace_id=trace_id,
                 message_id=message_id,
@@ -631,8 +671,7 @@ class GatewayClient:
                 metadata=metadata,
             )
             observation_id = getattr(langfuse_client_dispatch, "id", "")
-            if observation_id:
-                langfuse_parent_observation_id = observation_id
+            langfuse_parent_observation_id = observation_id or trace_parent_span_id
 
         header = MessageHeader(
             message_id=message_id,
@@ -734,6 +773,13 @@ class GatewayClient:
                         output={"success": False, "error": availability.error},
                         error=availability.error,
                     )
+                    if is_root_dispatch:
+                        await self._write_trace_root_end(
+                            trace_id=trace_id,
+                            status="FAILED",
+                            end_ts=int(time.time() * 1000),
+                            error=availability.error or "",
+                        )
                     return response
                 if availability.status == AvailabilityStatus.QUEUE_PENDING:
                     should_dispatch_control = False
@@ -750,6 +796,13 @@ class GatewayClient:
                 output={"success": False, "error": str(err)},
                 error=str(err),
             )
+            if is_root_dispatch:
+                await self._write_trace_root_end(
+                    trace_id=trace_id,
+                    status="FAILED",
+                    end_ts=int(time.time() * 1000),
+                    error=str(err),
+                )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -766,6 +819,13 @@ class GatewayClient:
                 output={"success": False, "error": str(err)},
                 error=str(err),
             )
+            if is_root_dispatch:
+                await self._write_trace_root_end(
+                    trace_id=trace_id,
+                    status="FAILED",
+                    end_ts=int(time.time() * 1000),
+                    error=str(err),
+                )
             return SendMessageResponse(
                 success=False,
                 status=ExecutionStatus.FAILED,
@@ -860,10 +920,10 @@ class GatewayClient:
         content: Any,
         metadata: Dict[str, Any],
     ) -> Any:
+        if self._langfuse_dispatch_fn is None:
+            return None
         try:
-            from by_framework_trace_langfuse import start_client_dispatch_observation
-
-            return start_client_dispatch_observation(
+            return self._langfuse_dispatch_fn(
                 trace_id=trace_id,
                 message_id=message_id,
                 target_agent_type=target_agent_type,
@@ -877,7 +937,6 @@ class GatewayClient:
             logger.warning(
                 "Langfuse client.dispatch observation skipped: %s",
                 err,
-                exc_info=True,
             )
             return None
 
@@ -951,3 +1010,52 @@ class GatewayClient:
             logger.warning(
                 "Failed to record client dispatch span: %s", err, exc_info=True
             )
+
+    async def _write_trace_root_start(
+        self,
+        *,
+        trace_id: str,
+        message_id: str,
+        session_id: str,
+        target_agent_type: str,
+        content: Any,
+        start_ts: int,
+    ) -> None:
+        """Write trace root start — best-effort, never propagates errors."""
+        try:
+            await self._trace_writer.record_trace(
+                TraceRecord(
+                    trace_id=trace_id,
+                    name=target_agent_type,
+                    session_id=session_id,
+                    root_message_id=message_id,
+                    root_agent_type=target_agent_type,
+                    input=sanitize_io_value(content, self._obs_config),
+                    status="QUEUED",
+                    start_ts=start_ts,
+                )
+            )
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass
+
+    async def _write_trace_root_end(
+        self,
+        *,
+        trace_id: str,
+        status: str,
+        end_ts: int,
+        error: str = "",
+    ) -> None:
+        """Write trace root end on routing failure — best-effort."""
+        try:
+            meta: dict = {"error": error} if error else {}
+            await self._trace_writer.record_trace(
+                TraceRecord(
+                    trace_id=trace_id,
+                    status=status,
+                    end_ts=end_ts,
+                    metadata=meta,
+                )
+            )
+        except Exception:  # pylint: disable=broad-exception-caught
+            pass

by_framework/common/config.py

@@ -50,6 +50,8 @@ class WorkerConfig:
     )
     heartbeat_lease_ttl_seconds: int = RedisKeys.WORKER_DEFAULT_LEASE_TTL_SECONDS
     lock_ttl_seconds: int = 60
+    worker_id_claim_max_wait_seconds: int = 90
+    worker_id_claim_retry_interval_seconds: float = 3.0
     stream_block_ms: int = 2000
 
 

by_framework/common/constants.py

@@ -136,7 +136,11 @@ class RedisKeys:
     # Set of known workers used for registry enumeration.
     KNOWN_WORKERS = "byai_gateway:registry:workers"
     WORKER_DEFAULT_HEARTBEAT_INTERVAL_SECONDS = 5
-    WORKER_DEFAULT_LEASE_TTL_SECONDS = 15
+    # 6× the heartbeat interval gives enough margin even when the main event
+    # loop is briefly occupied by an LLM call.  The dedicated heartbeat thread
+    # makes starvation practically impossible, but the larger TTL is a second
+    # line of defence and also gives Worker 2 time to claim after a crash.
+    WORKER_DEFAULT_LEASE_TTL_SECONDS = 30
 
     # Default TTL (7 days) for cleaning up session-related aggregate Keys
     DEFAULT_SESSION_TTL = 7 * 24 * 3600

by_framework/common/logger.py

@@ -167,7 +167,19 @@ def get_logger(name: str) -> logging.Logger:
     Returns:
         Logger instance
     """
-    return logging.getLogger(name)
+    base_logger = setup_logging()
+    if name == base_logger.name:
+        return base_logger
+
+    child_logger = logging.getLogger(name)
+    child_logger.setLevel(base_logger.level)
+    child_logger.propagate = False
+
+    if not child_logger.handlers:
+        for handler in base_logger.handlers:
+            child_logger.addHandler(handler)
+
+    return child_logger
 
 
 # Expose default logger

by_framework/core/__init__.py

@@ -0,0 +1,95 @@
+"""
+Core module of the Gateway SDK.
+
+Provides core components including protocol definitions, worker registry,
+and workspace management.
+"""
+
+from .availability import (
+    AvailabilityResult,
+    AvailabilityRouter,
+    AvailabilityStatus,
+    DeliveryIntent,
+    PendingDelivery,
+    RoutePolicy,
+    WakeupDecision,
+    WakeupDecisionStatus,
+    WakeupRequest,
+)
+from .delivery_gate import DeliveryGate
+from .protocol import (
+    ActionType,
+    AgentState,
+    ArtifactEvent,
+    AskAgentCommand,
+    AskUserEvent,
+    BaiYingMessage,
+    BaiYingMessageRole,
+    BaseCommand,
+    CancelTaskCommand,
+    CancelTaskResponse,
+    DataMessage,
+    EventType,
+    GatewayCommand,
+    MessageContent,
+    MessageFile,
+    MessageHeader,
+    Resource,
+    ResumeCommand,
+    SendMessageResponse,
+    SseMessageType,
+    SseReasonMessageType,
+    StateChangeEvent,
+    StreamChunkEvent,
+    command_from_dict,
+    get_registered_command,
+    register_command,
+    unregister_command,
+)
+from .registry import WorkerRegistry
+from .wakeup_controller import WakeupController, WakeupProvider
+from .workspace import WorkspaceManager
+
+__all__ = [
+    "SendMessageResponse",
+    "CancelTaskResponse",
+    "BaiYingMessage",
+    "BaiYingMessageRole",
+    "MessageContent",
+    "MessageFile",
+    "Resource",
+    "DataMessage",
+    "ActionType",
+    "AgentState",
+    "EventType",
+    "StateChangeEvent",
+    "StreamChunkEvent",
+    "ArtifactEvent",
+    "AskUserEvent",
+    "SseMessageType",
+    "SseReasonMessageType",
+    "MessageHeader",
+    "BaseCommand",
+    "AskAgentCommand",
+    "ResumeCommand",
+    "CancelTaskCommand",
+    "GatewayCommand",
+    "command_from_dict",
+    "register_command",
+    "unregister_command",
+    "get_registered_command",
+    "WorkerRegistry",
+    "WorkspaceManager",
+    "AvailabilityRouter",
+    "AvailabilityStatus",
+    "AvailabilityResult",
+    "DeliveryIntent",
+    "RoutePolicy",
+    "PendingDelivery",
+    "WakeupRequest",
+    "WakeupDecision",
+    "WakeupDecisionStatus",
+    "WakeupController",
+    "WakeupProvider",
+    "DeliveryGate",
+]