Fix: durable agent context (#223)

* feat: wip on orchestrator state fixing + tracing

Signed-off-by: Samantha Coyle <[email protected]>

* fix: separate ex/in-ternal triggers + wip fix orchestrators

Signed-off-by: Samantha Coyle <[email protected]>

* fix: ensure progress on substeps/steps

Signed-off-by: Samantha Coyle <[email protected]>

* fix: give orchestrators ability to pick up where they left off using same session id

Signed-off-by: Samantha Coyle <[email protected]>

* style: make linter happy

Signed-off-by: Samantha Coyle <[email protected]>

* fix: rm extra edge check since captured elsewhere

Signed-off-by: Samantha Coyle <[email protected]>

* feat: add session context for long term context on durable agents

Signed-off-by: Samantha Coyle <[email protected]>

* fix: address final merge conflict

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e ruff

Signed-off-by: Samantha Coyle <[email protected]>

* fix: rm local changes on requirements file for my testing

Signed-off-by: Samantha Coyle <[email protected]>

* style: tox -e flake8 fix

Signed-off-by: Samantha Coyle <[email protected]>

* fix: updates for tests too

Signed-off-by: Samantha Coyle <[email protected]>

* style: add todo comment

Signed-off-by: Samantha Coyle <[email protected]>

---------

Signed-off-by: Samantha Coyle <[email protected]>

Author: sicoyle

dapr_agents/agents/durableagent/agent.py

@@ -20,6 +20,7 @@
 from dapr_agents.types.workflow import DaprWorkflowStatus
 from dapr_agents.workflow.agentic import AgenticWorkflow
 from dapr_agents.workflow.decorators import message_router, task, workflow
+from dapr_agents.memory import ConversationDaprStateMemory
 
 from .schemas import (
     AgentTaskResponse,
@@ -59,6 +60,12 @@ class DurableAgent(AgenticWorkflow, AgentBase):
         default=None,
         description="The current workflow instance ID for this agent.",
     )
+    memory: ConversationDaprStateMemory = Field(
+        default_factory=lambda: ConversationDaprStateMemory(
+            store_name="workflowstatestore", session_id="durable_agent_session"
+        ),
+        description="Persistent memory with session-based state hydration.",
+    )
 
     @model_validator(mode="before")
     def set_agent_and_topic_name(cls, values: dict):
@@ -86,19 +93,23 @@ def model_post_init(self, __context: Any) -> None:
         if not self.state:
             self.state = {"instances": {}}
 
-        # Load the current workflow instance ID from state if it exists
+        # Load the current workflow instance ID from state using session_id
         logger.debug(f"State after loading: {self.state}")
         if self.state and self.state.get("instances"):
             logger.debug(f"Found {len(self.state['instances'])} instances in state")
             for instance_id, instance_data in self.state["instances"].items():
                 stored_workflow_name = instance_data.get("workflow_name")
+                stored_session_id = instance_data.get("session_id")
                 logger.debug(
-                    f"Instance {instance_id}: workflow_name={stored_workflow_name}, current_workflow_name={self._workflow_name}"
+                    f"Instance {instance_id}: workflow_name={stored_workflow_name}, session_id={stored_session_id}, current_workflow_name={self._workflow_name}, current_session_id={self.memory.session_id}"
                 )
-                if stored_workflow_name == self._workflow_name:
+                if (
+                    stored_workflow_name == self._workflow_name
+                    and stored_session_id == self.memory.session_id
+                ):
                     self.workflow_instance_id = instance_id
                     logger.debug(
-                        f"Loaded current workflow instance ID from state: {instance_id}"
+                        f"Loaded current workflow instance ID from state using session_id: {instance_id}"
                     )
                     break
         else:
@@ -256,9 +267,10 @@ def tool_calling_workflow(self, ctx: DaprWorkflowContext, message: TriggerAction
                                 "tool_call": tc,
                                 "instance_id": ctx.instance_id,
                                 "time": ctx.current_utc_datetime.isoformat(),
+                                "execution_order": i,  # Add ordering information
                             },
                         )
-                        for tc in tool_calls
+                        for i, tc in enumerate(tool_calls)
                     ]
                     yield self.when_all(parallel)
 
@@ -389,6 +401,7 @@ def record_initial_entry(
             "workflow_instance_id": instance_id,
             "triggering_workflow_instance_id": triggering_workflow_instance_id,
             "workflow_name": self._workflow_name,
+            "session_id": self.memory.session_id,
             "start_time": start_time_str,
             "trace_context": trace_context,
             "status": DaprWorkflowStatus.RUNNING.value,
@@ -430,6 +443,7 @@ def _ensure_instance_exists(
                 "workflow_instance_id": instance_id,
                 "triggering_workflow_instance_id": triggering_workflow_instance_id,
                 "workflow_name": self._workflow_name,
+                "session_id": self.memory.session_id,
                 "messages": [],
                 "tool_history": [],
                 "status": DaprWorkflowStatus.RUNNING.value,
@@ -581,6 +595,7 @@ def _create_tool_message_objects(self, tool_result: Dict[str, Any]) -> tuple:
             tool_call_id=tool_result["tool_call_id"],
             name=tool_result["tool_name"],
             content=tool_result["execution_result"],
+            role="tool",
         )
         agent_msg = DurableAgentMessage(**tool_msg.model_dump())
         tool_history_entry = ToolExecutionRecord(**tool_result)
@@ -656,7 +671,11 @@ def _get_last_message_from_state(
 
     @task
     async def run_tool(
-        self, tool_call: Dict[str, Any], instance_id: str, time: datetime
+        self,
+        tool_call: Dict[str, Any],
+        instance_id: str,
+        time: datetime,
+        execution_order: int = 0,
     ) -> Dict[str, Any]:
         """
         Executes a tool call atomically by invoking the specified function with the provided arguments
@@ -859,7 +878,7 @@ async def process_broadcast_message(self, message: BroadcastMessage):
 
                 # Start the agent's workflow
                 await self.run_and_monitor_workflow_async(
-                    workflow="ToolCallingWorkflow", input=trigger_message
+                    workflow="AgenticWorkflow", input=trigger_message
                 )
 
         except Exception as e:
@@ -871,9 +890,9 @@ def _construct_messages_with_instance_history(
         self, instance_id: str, input_data: Union[str, Dict[str, Any]]
     ) -> List[Dict[str, Any]]:
         """
-        Construct messages using instance-specific chat history instead of global memory.
-        This ensures proper message sequence for tool calls and prevents OpenAI API errors
-        in the event an app gets terminated or restarts while the workflow is running.
+        Construct messages using instance-specific chat history and persistent memory.
+        This ensures proper message sequence for tool calls and maintains conversation
+        history across workflow executions using the session_id.
 
         Args:
             instance_id: The workflow instance ID
@@ -882,12 +901,13 @@ def _construct_messages_with_instance_history(
         Returns:
             List of formatted messages with proper sequence
         """
+        additional_context_messages: List[Dict[str, Any]] = []
         if not self.prompt_template:
             raise ValueError(
                 "Prompt template must be initialized before constructing messages."
             )
 
-        # Get instance-specific chat history instead of global memory
+        # Get instance-specific chat history
         if self.state is None:
             logger.warning(
                 f"Agent state is None for instance {instance_id}, initializing empty state"
@@ -900,48 +920,51 @@ def _construct_messages_with_instance_history(
         else:
             instance_messages = []
 
-        # Always include long-term memory (chat_history) for context
-        # This ensures agents have access to broadcast messages and persistent context
-        long_term_memory_data = self.state.get("chat_history", [])
+        # Get messages from persistent memory (session-based, cross-workflow)
+        persistent_memory_messages = []
+        try:
+            persistent_memory_messages = self.memory.get_messages()
+            logger.info(
+                f"Retrieved {len(persistent_memory_messages)} messages for session {self.memory.session_id}"
+            )
+        except Exception as e:
+            logger.warning(f"Failed to retrieve persistent memory: {e}")
 
-        # Convert long-term memory to dict format for LLM consumption
+        # Get long-term memory from workflow state (for broadcast messages and persistent context)
+        long_term_memory_data = self.state.get("chat_history", [])
         long_term_memory_messages = []
         for msg in long_term_memory_data:
             if isinstance(msg, dict):
                 long_term_memory_messages.append(msg)
             elif hasattr(msg, "model_dump"):
                 long_term_memory_messages.append(msg.model_dump())
 
-        # For broadcast-triggered workflows, also include additional context memory
-        source = instance_data.get("source") if instance_data else None
-        additional_context_messages = []
-        if source and source != "direct":
-            # Include additional context memory for broadcast-triggered workflows
-            context_memory_data = self.memory.get_messages()
-            for msg in context_memory_data:
-                if isinstance(msg, dict):
-                    additional_context_messages.append(msg)
-                elif hasattr(msg, "model_dump"):
-                    additional_context_messages.append(msg.model_dump())
-
-        # Build chat history with:
-        # 1. Long-term memory (persistent context, broadcast messages)
-        # 2. Short-term instance messages (current workflow specific)
-        # 3. Additional context memory (for broadcast-triggered workflows)
+        # Build chat history with proper context and order
         chat_history = []
 
-        # Add long-term memory first (broadcast messages, persistent context)
-        chat_history.extend(long_term_memory_messages)
+        # First add persistent memory and long-term memory as user messages for context
+        # This ensures we have cross-workflow context but doesn't interfere with tool state order
+        for msg in persistent_memory_messages + long_term_memory_messages:
+            msg_dict = msg.model_dump() if hasattr(msg, "model_dump") else dict(msg)
+            if msg_dict in chat_history:
+                continue
+            # TODO: We need to properly design session-based memory.
+            # Convert tool-related messages to user messages to avoid conversation order issues
+            if msg_dict.get("role") in ["tool", "assistant"] and (
+                msg_dict.get("tool_calls") or msg_dict.get("tool_call_id")
+            ):
+                msg_dict = {
+                    "role": "user",
+                    "content": f"[Previous {msg_dict['role']} message: {msg_dict.get('content', '')}]",
+                }
+            chat_history.append(msg_dict)
 
-        # Add short-term instance messages (current workflow)
+        # Then add instance messages in their original form to maintain tool state
         for msg in instance_messages:
-            if isinstance(msg, dict):
-                chat_history.append(msg)
-            else:
-                # Convert DurableAgentMessage to dict if needed
-                chat_history.append(
-                    msg.model_dump() if hasattr(msg, "model_dump") else dict(msg)
-                )
+            msg_dict = msg.model_dump() if hasattr(msg, "model_dump") else dict(msg)
+            if msg_dict in chat_history:
+                continue
+            chat_history.append(msg_dict)
 
         # Add additional context memory last (for broadcast-triggered workflows)
         chat_history.extend(additional_context_messages)

quickstarts/03-agent-tool-call/requirements.txt

@@ -3,4 +3,4 @@ arize-phoenix>=11.22.0,<12.0.0
 arize-phoenix-otel>=0.12.0,<0.13.0
 opentelemetry-instrumentation>=0.56b0
 # For local development use local changes by commenting out the dapr-agents line above and uncommenting the line below:
-# -e ../../
+# -e ../../[observability]

quickstarts/05-multi-agent-workflows/requirements.txt

@@ -3,4 +3,4 @@ arize-phoenix>=11.22.0,<12.0.0
 arize-phoenix-otel>=0.12.0,<0.13.0
 opentelemetry-instrumentation>=0.56b0
 # For local development use local changes by commenting out the dapr-agents line above and uncommenting the line below:
-# -e ../../[observability[
+# -e ../../[observability]

tests/agents/durableagent/test_durable_agent.py

@@ -5,7 +5,7 @@
 import asyncio
 import os
 from typing import Any
-from unittest.mock import AsyncMock, Mock, patch
+from unittest.mock import AsyncMock, Mock, patch, MagicMock
 
 import pytest
 from dapr.ext.workflow import DaprWorkflowContext
@@ -21,7 +21,7 @@
     DurableAgentWorkflowState,
 )
 from dapr_agents.llm import OpenAIChatClient
-from dapr_agents.memory import ConversationListMemory
+from dapr_agents.memory import ConversationDaprStateMemory
 from dapr_agents.tool.base import AgentTool
 from dapr_agents.types import (
     AssistantMessage,
@@ -104,13 +104,43 @@ def mock_register_agentic_system(self):
     yield
 
 
+class MockDaprClient:
+    """Mock DaprClient that supports context manager protocol"""
+
+    def __init__(self):
+        self.get_state = MagicMock(return_value=Mock(data=None, json=lambda: {}))
+        self.save_state = MagicMock()
+        self.delete_state = MagicMock()
+        self.query_state = MagicMock()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def __call__(self, *args, **kwargs):
+        return self
+
+
 class TestDurableAgent:
     """Test cases for the DurableAgent class."""
 
     @pytest.fixture(autouse=True)
-    def setup_env(self):
-        """Set up environment variables for testing."""
+    def setup_env(self, monkeypatch):
+        """Set up environment variables and mocks for testing."""
         os.environ["OPENAI_API_KEY"] = "test-api-key"
+
+        # Mock DaprClient to use our context manager supporting mock
+        mock_client = MockDaprClient()
+        mock_client.get_state.return_value = Mock(data=None)  # Default empty state
+
+        # Patch both the client import locations
+        monkeypatch.setattr("dapr.clients.DaprClient", lambda: mock_client)
+        monkeypatch.setattr(
+            "dapr_agents.storage.daprstores.statestore.DaprClient", lambda: mock_client
+        )
+
         yield
         if "OPENAI_API_KEY" in os.environ:
             del os.environ["OPENAI_API_KEY"]
@@ -158,7 +188,9 @@ def basic_durable_agent(self, mock_llm):
             goal="Help with testing",
             instructions=["Be helpful", "Test things"],
             llm=mock_llm,
-            memory=ConversationListMemory(),
+            memory=ConversationDaprStateMemory(
+                store_name="teststatestore", session_id="test_session"
+            ),
             max_iterations=5,
             state_store_name="teststatestore",
             message_bus_name="testpubsub",
@@ -174,7 +206,9 @@ def durable_agent_with_tools(self, mock_llm, mock_tool):
             goal="Execute tools",
             instructions=["Use tools when needed"],
             llm=mock_llm,
-            memory=ConversationListMemory(),
+            memory=ConversationDaprStateMemory(
+                store_name="teststatestore", session_id="test_session"
+            ),
             tools=[mock_tool],
             max_iterations=5,
             state_store_name="teststatestore",