sdk/python: Bug fixes for lifecycle hooks

Author: ndeodhar

sdk/python/polos/agents/stream.py

@@ -157,8 +157,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
         if agent and agent.on_agent_step_start:
             hook_context = HookContext(
                 workflow_id=ctx.agent_id,
-                agent_workflow_id=ctx.agent_id,
-                agent_run_id=agent_run_id,
                 session_id=ctx.session_id,
                 user_id=ctx.user_id,
                 agent_config=agent_config,
@@ -173,8 +171,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
             )
 
             # Apply modifications
-            if hook_result.modified_agent_config:
-                agent_config.update(hook_result.modified_agent_config)
             if hook_result.modified_payload and "messages" in hook_result.modified_payload:
                 conversation_messages = hook_result.modified_payload["messages"]
 
@@ -284,8 +280,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
             if agent and agent.on_tool_start:
                 hook_context = HookContext(
                     workflow_id=ctx.agent_id,
-                    agent_workflow_id=ctx.agent_id,
-                    agent_run_id=agent_run_id,
                     session_id=ctx.session_id,
                     user_id=ctx.user_id,
                     agent_config=agent_config,
@@ -298,8 +292,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
                 )
 
                 # Apply modifications
-                if hook_result.modified_agent_config:
-                    agent_config.update(hook_result.modified_agent_config)
                 if hook_result.modified_payload:
                     tool_args.update(hook_result.modified_payload)
 
@@ -353,8 +345,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
                 if agent and agent.on_tool_end:
                     hook_context = HookContext(
                         workflow_id=ctx.agent_id,
-                        agent_workflow_id=ctx.agent_id,
-                        agent_run_id=agent_run_id,
                         session_id=ctx.session_id,
                         user_id=ctx.user_id,
                         agent_config=agent_config,
@@ -368,8 +358,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
                     )
 
                     # Apply modifications
-                    if hook_result.modified_agent_config:
-                        agent_config.update(hook_result.modified_agent_config)
                     if hook_result.modified_output is not None:
                         tool_result = hook_result.modified_output
 
@@ -442,8 +430,6 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
         if agent and agent.on_agent_step_end:
             hook_context = HookContext(
                 workflow_id=ctx.agent_id,
-                agent_workflow_id=ctx.agent_id,
-                agent_run_id=agent_run_id,
                 session_id=ctx.session_id,
                 user_id=ctx.user_id,
                 agent_config=agent_config,
@@ -456,11 +442,9 @@ async def _agent_stream_function(ctx: AgentContext, payload: dict[str, Any]) ->
             )
 
             # Apply modifications
-            if hook_result.modified_agent_config:
-                agent_config.update(hook_result.modified_agent_config)
             if hook_result.modified_output:
                 new_result = hook_result.modified_output
-                steps[-1].update(new_result)
+                steps[-1] = new_result
 
             # Check hook action
             if hook_result.action == HookAction.FAIL:

sdk/python/polos/core/step.py

@@ -188,7 +188,7 @@ async def _publish_step_event(
                 data={
                     "step_key": step_key,
                     "step_type": step_type,
-                    "input_params": safe_serialize(input_params) if input_params else {},
+                    "data": safe_serialize(input_params) if input_params else {},
                     "_metadata": {
                         "execution_id": self.ctx.execution_id,
                         "workflow_id": self.ctx.workflow_id,

sdk/python/polos/middleware/guardrail.py

@@ -26,8 +26,8 @@ class GuardrailContext(BaseModel):
     tool_calls: list[ToolCall] | None = None  # LLM tool calls
 
     # Execution context (for guardrail to make decisions)
-    agent_workflow_id: str = ""
-    agent_run_id: str = ""
+    agent_workflow_id: str | None = ""
+    agent_run_id: str | None = ""
     session_id: str | None = None
     user_id: str | None = None
     llm_config: AgentConfig = AgentConfig(name="", provider="", model="")

sdk/python/polos/middleware/hook.py

@@ -32,19 +32,17 @@ class HookContext(BaseModel):
 
     # Immutable identifiers
     workflow_id: str
-    agent_workflow_id: str | None = None  # Available for agents
-    agent_run_id: str | None = None  # Available for agents
     session_id: str | None = None
     user_id: str | None = None
-    agent_config: AgentConfig | None = None
+    agent_config: AgentConfig | None = None  # Not available to agent's on_start and on_end hooks
 
     # Current state
     steps: list[Step] = []  # All previous steps
 
     # For workflow/tool hooks
     current_tool: str | None = None
-    current_payload: dict[str, Any] | None = None
-    current_output: dict[str, Any] | None = None
+    current_payload: dict[str, Any] | BaseModel | None = None
+    current_output: dict[str, Any] | BaseModel | None = None
 
     def to_dict(self) -> dict[str, Any]:
         """Convert hook context to dictionary for serialization."""
@@ -72,7 +70,6 @@ class HookResult(BaseModel):
     action: HookAction = HookAction.CONTINUE
 
     # Optional modifications
-    modified_agent_config: AgentConfig | None = None
     modified_payload: dict[str, Any] | None = None
     modified_output: Any | None = None
 
@@ -84,7 +81,7 @@ def continue_with(cls, **modifications) -> "HookResult":
         """Continue with optional modifications.
 
         Args:
-            **modifications: Can include modified_agent_config, modified_payload, modified_output
+            **modifications: Can include modified_payload, modified_output
 
         Returns:
             HookResult with CONTINUE action and modifications

sdk/python/polos/middleware/hooks.py

@@ -243,7 +243,7 @@ async def test_publish_step_event(self, mock_workflow_context):
             assert event.event_type == event_type
             assert event.data["step_key"] == step_key
             assert event.data["step_type"] == event_name  # Note: it's "step_type" not "event_name"
-            assert "input_params" in event.data
+            assert "data" in event.data
 
     @pytest.mark.asyncio
     async def test_publish_step_event_topic(self, mock_workflow_context):

sdk/python/tests/unit/test_core/test_step.py

@@ -28,8 +28,6 @@ def test_hook_context_initialization(self):
         """Test HookContext initialization."""
         ctx = HookContext(workflow_id="test-workflow")
         assert ctx.workflow_id == "test-workflow"
-        assert ctx.agent_workflow_id is None
-        assert ctx.agent_run_id is None
         assert ctx.session_id is None
         assert ctx.user_id is None
         assert ctx.agent_config is None
@@ -46,8 +44,6 @@ def test_hook_context_full_initialization(self):
         steps = [Step(step=1, content="test")]
         ctx = HookContext(
             workflow_id="test-workflow",
-            agent_workflow_id="test-agent",
-            agent_run_id="test-run",
             session_id="test-session",
             user_id="test-user",
             agent_config=agent_config,
@@ -57,8 +53,6 @@ def test_hook_context_full_initialization(self):
             current_output={"result": "output"},
         )
         assert ctx.workflow_id == "test-workflow"
-        assert ctx.agent_workflow_id == "test-agent"
-        assert ctx.agent_run_id == "test-run"
         assert ctx.session_id == "test-session"
         assert ctx.user_id == "test-user"
         assert ctx.agent_config == agent_config
@@ -106,7 +100,6 @@ def test_hook_result_default(self):
         """Test HookResult with default values."""
         result = HookResult()
         assert result.action == HookAction.CONTINUE
-        assert result.modified_agent_config is None
         assert result.modified_payload is None
         assert result.modified_output is None
         assert result.error_message is None