agnt5 0.2.2__cp39-abi3-macosx_11_0_arm64.whl → 0.2.4__cp39-abi3-macosx_11_0_arm64.whl

agnt5/worker.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import asyncio
+import contextvars
 import logging
 from typing import Any, Dict, Optional
 
@@ -12,6 +13,12 @@ from ._telemetry import setup_module_logger
 
 logger = setup_module_logger(__name__)
 
+# Context variable to store trace metadata for propagation to LM calls
+# This allows Rust LM layer to access traceparent without explicit parameter passing
+_trace_metadata: contextvars.ContextVar[Dict[str, str]] = contextvars.ContextVar(
+    '_trace_metadata', default={}
+)
+
 
 class Worker:
     """AGNT5 Worker for registering and running functions/workflows with the coordinator.
@@ -88,6 +95,10 @@ class Worker:
         # Create Rust worker instance
         self._rust_worker = self._PyWorker(self._rust_config)
 
+        # Create worker-scoped entity state manager
+        from .entity import EntityStateManager
+        self._entity_state_manager = EntityStateManager()
+
         logger.info(
             f"Worker initialized: {service_name} v{service_version} (runtime: {runtime})"
         )
@@ -191,7 +202,7 @@ class Worker:
 
             # Build metadata dict with methods list and schemas
             metadata_dict = {
-                "methods": json.dumps(list(entity_type._methods.keys())),
+                "methods": json.dumps(list(entity_type._method_schemas.keys())),
                 "method_schemas": json.dumps(method_schemas)
             }
 
@@ -205,7 +216,7 @@ class Worker:
                 definition=None,
             )
             components.append(component_info)
-            logger.debug(f"Discovered entity: {name} with methods: {list(entity_type._methods.keys())}")
+            logger.debug(f"Discovered entity: {name} with methods: {list(entity_type._method_schemas.keys())}")
 
         # Discover agents
         for name, agent in AgentRegistry.all().items():
@@ -236,68 +247,6 @@ class Worker:
             components.append(component_info)
             logger.debug(f"Discovered agent: {name}")
 
-        # Discover tools
-        for name, tool in ToolRegistry.all().items():
-            # Serialize schemas to JSON strings
-            input_schema_str = None
-            if hasattr(tool, 'input_schema') and tool.input_schema:
-                input_schema_str = json.dumps(tool.input_schema)
-
-            output_schema_str = None
-            if hasattr(tool, 'output_schema') and tool.output_schema:
-                output_schema_str = json.dumps(tool.output_schema)
-
-            component_info = self._PyComponentInfo(
-                name=name,
-                component_type="tool",
-                metadata={},
-                config={},
-                input_schema=input_schema_str,
-                output_schema=output_schema_str,
-                definition=None,
-            )
-            components.append(component_info)
-            logger.debug(f"Discovered tool: {name}")
-
-        # Discover entities
-        for name, entity_type in EntityRegistry.all().items():
-            # Build metadata dict with methods list as JSON string
-            metadata_dict = {
-                "methods": json.dumps(list(entity_type._methods.keys()))
-            }
-
-            component_info = self._PyComponentInfo(
-                name=name,
-                component_type="entity",
-                metadata=metadata_dict,
-                config={},
-                input_schema=None,
-                output_schema=None,
-                definition=None,
-            )
-            components.append(component_info)
-            logger.debug(f"Discovered entity: {name} with methods: {list(entity_type._methods.keys())}")
-
-        # Discover agents
-        for name, agent in AgentRegistry.all().items():
-            # Build metadata dict with agent info
-            metadata_dict = {
-                "model": agent.model_name,
-                "tools": json.dumps(list(agent.tools.keys()) if hasattr(agent, 'tools') else [])
-            }
-
-            component_info = self._PyComponentInfo(
-                name=name,
-                component_type="agent",
-                metadata=metadata_dict,
-                config={},
-                input_schema=None,
-                output_schema=None,
-                definition=None,
-            )
-            components.append(component_info)
-            logger.debug(f"Discovered agent: {name}")
-
         logger.info(f"Discovered {len(components)} components")
         return components
 
@@ -305,67 +254,66 @@ class Worker:
         """Create the message handler that will be called by Rust worker."""
 
         def handle_message(request):
-            """Handle incoming execution requests."""
-            try:
-                # Extract request details
-                component_name = request.component_name
-                component_type = request.component_type
-                input_data = request.input_data
-
-                logger.debug(
-                    f"Handling {component_type} request: {component_name}, input size: {len(input_data)} bytes"
-                )
+            """Handle incoming execution requests - returns coroutine for Rust to await."""
+            # Extract request details
+            component_name = request.component_name
+            component_type = request.component_type
+            input_data = request.input_data
+
+            logger.debug(
+                f"Handling {component_type} request: {component_name}, input size: {len(input_data)} bytes"
+            )
 
-                # Import all registries
-                from .tool import ToolRegistry
-                from .entity import EntityRegistry
-                from .agent import AgentRegistry
-
-                # Route based on component type
-                if component_type == "tool":
-                    tool = ToolRegistry.get(component_name)
-                    if tool:
-                        logger.debug(f"Found tool: {component_name}")
-                        result = asyncio.run(self._execute_tool(tool, input_data, request))
-                        return result
-
-                elif component_type == "entity":
-                    entity_type = EntityRegistry.get(component_name)
-                    if entity_type:
-                        logger.debug(f"Found entity: {component_name}")
-                        result = asyncio.run(self._execute_entity(entity_type, input_data, request))
-                        return result
-
-                elif component_type == "agent":
-                    agent = AgentRegistry.get(component_name)
-                    if agent:
-                        logger.debug(f"Found agent: {component_name}")
-                        result = asyncio.run(self._execute_agent(agent, input_data, request))
-                        return result
-
-                elif component_type == "workflow":
-                    workflow_config = WorkflowRegistry.get(component_name)
-                    if workflow_config:
-                        logger.debug(f"Found workflow: {component_name}")
-                        result = asyncio.run(self._execute_workflow(workflow_config, input_data, request))
-                        return result
-
-                elif component_type == "function":
-                    function_config = FunctionRegistry.get(component_name)
-                    if function_config:
-                        logger.info(f"🔥 WORKER: Received request for function: {component_name}")
-                        result = asyncio.run(self._execute_function(function_config, input_data, request))
-                        return result
-
-                # Not found
-                error_msg = f"Component '{component_name}' of type '{component_type}' not found"
-                logger.error(error_msg)
+            # Import all registries
+            from .tool import ToolRegistry
+            from .entity import EntityRegistry
+            from .agent import AgentRegistry
+
+            # Route based on component type and return coroutines
+            if component_type == "tool":
+                tool = ToolRegistry.get(component_name)
+                if tool:
+                    logger.debug(f"Found tool: {component_name}")
+                    # Return coroutine, don't await it
+                    return self._execute_tool(tool, input_data, request)
+
+            elif component_type == "entity":
+                entity_type = EntityRegistry.get(component_name)
+                if entity_type:
+                    logger.debug(f"Found entity: {component_name}")
+                    # Return coroutine, don't await it
+                    return self._execute_entity(entity_type, input_data, request)
+
+            elif component_type == "agent":
+                agent = AgentRegistry.get(component_name)
+                if agent:
+                    logger.debug(f"Found agent: {component_name}")
+                    # Return coroutine, don't await it
+                    return self._execute_agent(agent, input_data, request)
+
+            elif component_type == "workflow":
+                workflow_config = WorkflowRegistry.get(component_name)
+                if workflow_config:
+                    logger.debug(f"Found workflow: {component_name}")
+                    # Return coroutine, don't await it
+                    return self._execute_workflow(workflow_config, input_data, request)
+
+            elif component_type == "function":
+                function_config = FunctionRegistry.get(component_name)
+                if function_config:
+                    logger.info(f"🔥 WORKER: Received request for function: {component_name}")
+                    # Return coroutine, don't await it
+                    return self._execute_function(function_config, input_data, request)
+
+            # Not found - need to return an async error response
+            error_msg = f"Component '{component_name}' of type '{component_type}' not found"
+            logger.error(error_msg)
+
+            # Create async wrapper for error response
+            async def error_response():
                 return self._create_error_response(request, error_msg)
 
-            except Exception as e:
-                error_msg = f"Handler error: {e}"
-                logger.error(error_msg, exc_info=True)
-                return self._create_error_response(request, error_msg)
+            return error_response()
 
         return handle_message
 
@@ -382,20 +330,45 @@ class Worker:
             # Parse input data
             input_dict = json.loads(input_data.decode("utf-8")) if input_data else {}
 
-            # Create context
+            # Store trace metadata in contextvar for LM calls to access
+            # The Rust worker injects traceparent into request.metadata for trace propagation
+            if hasattr(request, 'metadata') and request.metadata:
+                _trace_metadata.set(dict(request.metadata))
+                logger.debug(f"Trace metadata stored: traceparent={request.metadata.get('traceparent', 'N/A')}")
+
+            # Create context with runtime_context for trace correlation
             ctx = Context(
                 run_id=f"{self.service_name}:{config.name}",
-                component_type="function",
+                runtime_context=request.runtime_context,
             )
 
-            # Execute function
-            if input_dict:
-                result = config.handler(ctx, **input_dict)
-            else:
-                result = config.handler(ctx)
-
-            # Debug: Log what type result is
-            logger.info(f"🔥 WORKER: Function result type: {type(result).__name__}, isasyncgen: {inspect.isasyncgen(result)}, iscoroutine: {inspect.iscoroutine(result)}")
+            # Create span for function execution with trace linking
+            from ._core import create_span, flush_telemetry_py
+
+            with create_span(
+                config.name,
+                "function",
+                request.runtime_context,
+                {
+                    "function.name": config.name,
+                    "service.name": self.service_name,
+                },
+            ) as span:
+                # Execute function
+                if input_dict:
+                    result = config.handler(ctx, **input_dict)
+                else:
+                    result = config.handler(ctx)
+
+                # Debug: Log what type result is
+                logger.info(f"🔥 WORKER: Function result type: {type(result).__name__}, isasyncgen: {inspect.isasyncgen(result)}, iscoroutine: {inspect.iscoroutine(result)}")
+
+            # Flush telemetry after span ends to ensure it's exported
+            try:
+                flush_telemetry_py()
+                logger.debug("Telemetry flushed after function execution")
+            except Exception as e:
+                logger.warning(f"Failed to flush telemetry: {e}")
 
             # Check if result is an async generator (streaming function)
             if inspect.isasyncgen(result):
@@ -473,16 +446,17 @@ class Worker:
             )
 
     async def _execute_workflow(self, config, input_data: bytes, request):
-        """Execute a workflow handler with replay support (Phase 6B)."""
+        """Execute a workflow handler with automatic replay support."""
         import json
-        from .context import Context
+        from .workflow import WorkflowEntity, WorkflowContext
+        from .entity import _get_state_manager
         from ._core import PyExecuteComponentResponse
 
         try:
             # Parse input data
             input_dict = json.loads(input_data.decode("utf-8")) if input_data else {}
 
-            # Phase 6B: Parse replay data from request metadata
+            # Parse replay data from request metadata for crash recovery
             completed_steps = {}
             initial_state = {}
 
@@ -507,38 +481,74 @@ class Worker:
                         except json.JSONDecodeError:
                             logger.warning("Failed to parse workflow_state from metadata")
 
-            # Create context with replay data
-            ctx = Context(
+            # Create WorkflowEntity for state management
+            workflow_entity = WorkflowEntity(run_id=f"{self.service_name}:{config.name}")
+
+            # Load replay data into entity if provided
+            if completed_steps:
+                workflow_entity._completed_steps = completed_steps
+                logger.debug(f"Loaded {len(completed_steps)} completed steps into workflow entity")
+
+            if initial_state:
+                # Load initial state into entity's state manager
+                state_manager = _get_state_manager()
+                state_manager._states[workflow_entity._state_key] = initial_state
+                logger.debug(f"Loaded initial state with {len(initial_state)} keys into workflow entity")
+
+            # Create WorkflowContext with entity and runtime_context for trace correlation
+            ctx = WorkflowContext(
+                workflow_entity=workflow_entity,
                 run_id=f"{self.service_name}:{config.name}",
-                component_type="workflow",
-                completed_steps=completed_steps if completed_steps else None,
-                initial_state=initial_state if initial_state else None,
+                runtime_context=request.runtime_context,
             )
 
-            # Execute workflow
-            if input_dict:
-                result = await config.handler(ctx, **input_dict)
-            else:
-                result = await config.handler(ctx)
+            # Create span for workflow execution with trace linking
+            from ._core import create_span, flush_telemetry_py
+
+            with create_span(
+                config.name,
+                "workflow",
+                request.runtime_context,
+                {
+                    "workflow.name": config.name,
+                    "service.name": self.service_name,
+                },
+            ) as span:
+                # Execute workflow
+                if input_dict:
+                    result = await config.handler(ctx, **input_dict)
+                else:
+                    result = await config.handler(ctx)
+
+            # Flush telemetry after span ends to ensure it's exported
+            try:
+                flush_telemetry_py()
+                logger.debug("Telemetry flushed after workflow execution")
+            except Exception as e:
+                logger.warning(f"Failed to flush telemetry: {e}")
 
             # Serialize result
             output_data = json.dumps(result).encode("utf-8")
 
-            # Phase 6: Collect workflow execution metadata (similar to entity pattern)
+            # Collect workflow execution metadata for durability
             metadata = {}
 
             # Add step events to metadata (for workflow durability)
-            if ctx._step_events:
-                metadata["step_events"] = json.dumps(ctx._step_events)
-                logger.debug(f"Workflow has {len(ctx._step_events)} recorded steps")
+            # Access _step_events from the workflow entity, not the context
+            step_events = ctx._workflow_entity._step_events
+            if step_events:
+                metadata["step_events"] = json.dumps(step_events)
+                logger.debug(f"Workflow has {len(step_events)} recorded steps")
 
             # Add final state snapshot to metadata (if state was used)
-            if hasattr(ctx, '_state_client') and ctx.state.has_changes():
-                state_snapshot = ctx.state.get_state_snapshot()
-                metadata["workflow_state"] = json.dumps(state_snapshot)
-                logger.debug(f"Workflow state snapshot: {state_snapshot}")
+            # Check if _state was initialized without triggering property getter
+            if hasattr(ctx, '_workflow_entity') and ctx._workflow_entity._state is not None:
+                if ctx._workflow_entity._state.has_changes():
+                    state_snapshot = ctx._workflow_entity._state.get_state_snapshot()
+                    metadata["workflow_state"] = json.dumps(state_snapshot)
+                    logger.debug(f"Workflow state snapshot: {state_snapshot}")
 
-            logger.info(f"Workflow completed successfully with {len(ctx._step_events)} steps")
+            logger.info(f"Workflow completed successfully with {len(step_events)} steps")
 
             return PyExecuteComponentResponse(
                 invocation_id=request.invocation_id,
@@ -578,10 +588,10 @@ class Worker:
             # Parse input data
             input_dict = json.loads(input_data.decode("utf-8")) if input_data else {}
 
-            # Create context
+            # Create context with runtime_context for trace correlation
             ctx = Context(
                 run_id=f"{self.service_name}:{tool.name}",
-                component_type="tool",
+                runtime_context=request.runtime_context,
             )
 
             # Execute tool
@@ -622,9 +632,12 @@ class Worker:
         """Execute an entity method."""
         import json
         from .context import Context
-        from .entity import EntityType, DurableEntity, _entity_states
+        from .entity import EntityType, Entity, _entity_state_manager_ctx
         from ._core import PyExecuteComponentResponse
 
+        # Set entity state manager in context for Entity instances to access
+        _entity_state_manager_ctx.set(self._entity_state_manager)
+
         try:
             # Parse input data
             input_dict = json.loads(input_data.decode("utf-8")) if input_data else {}
@@ -638,13 +651,26 @@ class Worker:
             if not method_name:
                 raise ValueError("Entity invocation requires 'method' parameter")
 
-            # Check if this is a class-based entity (DurableEntity subclass) or method-based (EntityType)
-            if entity_type.entity_class is not None:
-                # Class-based entity: use the stored class reference
-                entity_instance = entity_type.entity_class(key=entity_key)
-            else:
-                # Method-based entity: create EntityInstance
-                entity_instance = entity_type(entity_key)
+            # Load state from platform if provided in request metadata
+            state_key = (entity_type.name, entity_key)
+            if hasattr(request, 'metadata') and request.metadata:
+                if "entity_state" in request.metadata:
+                    platform_state_json = request.metadata["entity_state"]
+                    platform_version = int(request.metadata.get("state_version", "0"))
+
+                    # Load platform state into state manager
+                    self._entity_state_manager.load_state_from_platform(
+                        state_key,
+                        platform_state_json,
+                        platform_version
+                    )
+                    logger.info(
+                        f"Loaded entity state from platform: {entity_type.name}/{entity_key} "
+                        f"(version {platform_version})"
+                    )
+
+            # Create entity instance using the stored class reference
+            entity_instance = entity_type.entity_class(key=entity_key)
 
             # Get method
             if not hasattr(entity_instance, method_name):
@@ -658,26 +684,32 @@ class Worker:
             # Serialize result
             output_data = json.dumps(result).encode("utf-8")
 
-            # Phase 5B: Capture entity state after execution for persistence
-            state_key = (entity_type.name, entity_key)
+            # Capture entity state after execution with version tracking
+            state_dict, expected_version, new_version = \
+                self._entity_state_manager.get_state_for_persistence(state_key)
+
             metadata = {}
-            if state_key in _entity_states:
-                entity_state = _entity_states[state_key]
+            if state_dict:
                 # Serialize state as JSON string for platform persistence
-                state_json = json.dumps(entity_state)
+                state_json = json.dumps(state_dict)
                 # Pass in metadata for Worker Coordinator to publish
                 metadata = {
                     "entity_state": state_json,
                     "entity_type": entity_type.name,
                     "entity_key": entity_key,
+                    "expected_version": str(expected_version),
+                    "new_version": str(new_version),
                 }
-                logger.debug(f"Entity state update: {entity_type.name}:{entity_key}, state: {state_json}")
+                logger.info(
+                    f"Captured entity state: {entity_type.name}/{entity_key} "
+                    f"(version {expected_version} → {new_version})"
+                )
 
             return PyExecuteComponentResponse(
                 invocation_id=request.invocation_id,
                 success=True,
                 output_data=output_data,
-                state_update=None,  # TODO: Phase 6 - Use structured StateUpdate object
+                state_update=None,  # TODO: Use structured StateUpdate object
                 error_message=None,
                 metadata=metadata,  # Include state in metadata for Worker Coordinator
                 is_chunk=False,
@@ -716,10 +748,10 @@ class Worker:
             if not user_message:
                 raise ValueError("Agent invocation requires 'message' parameter")
 
-            # Create context
+            # Create context with runtime_context for trace correlation
             ctx = Context(
                 run_id=f"{self.service_name}:{agent.name}",
-                component_type="agent",
+                runtime_context=request.runtime_context,
             )
 
             # Execute agent