From ff698271a5e0afa74f24fa9b3979545e97afbabe Mon Sep 17 00:00:00 2001 From: Anatoli Beliaev Date: Thu, 11 Sep 2025 13:36:13 -0700 Subject: [PATCH 1/5] Make DurableAIAgentContext delegate to the underlying DurableOrchestrationContext automatically --- .../openai_agents/context.py | 55 +++- tests/openai_agents/test_context.py | 292 ++++++++++++++++++ 2 files changed, 336 insertions(+), 11 deletions(-) create mode 100644 tests/openai_agents/test_context.py diff --git a/azure/durable_functions/openai_agents/context.py b/azure/durable_functions/openai_agents/context.py index 973ab01d..72657344 100644 --- a/azure/durable_functions/openai_agents/context.py +++ b/azure/durable_functions/openai_agents/context.py @@ -1,4 +1,4 @@ -from typing import Any, Callable, Optional +from typing import Any, Callable, Optional, TYPE_CHECKING from azure.durable_functions.models.DurableOrchestrationContext import ( DurableOrchestrationContext, @@ -11,8 +11,38 @@ from .task_tracker import TaskTracker -class DurableAIAgentContext: - """Context for AI agents running in Azure Durable Functions orchestration.""" +if TYPE_CHECKING: + # At type-check time we want all members / signatures for IDE & linters. + _BaseDurableContext = DurableOrchestrationContext +else: + class _BaseDurableContext: # lightweight runtime stub + """Runtime stub base class for delegation; real context is wrapped. + + At runtime we avoid inheriting from DurableOrchestrationContext so that + attribute lookups for its members are delegated via __getattr__ to the + wrapped ``_context`` instance. + """ + __slots__ = () + + +class DurableAIAgentContext(_BaseDurableContext): + """Context for AI agents running in Azure Durable Functions orchestration. + + Design + ------ + * Static analysis / IDEs: Appears to subclass ``DurableOrchestrationContext`` so + you get autocompletion and type hints (under TYPE_CHECKING branch). + * Runtime: Inherits from a trivial stub. All durable orchestration operations + are delegated to the real ``DurableOrchestrationContext`` instance provided + as ``context`` and stored in ``_context``. + + Consequences + ------------ + * ``isinstance(DurableAIAgentContext, DurableOrchestrationContext)`` is **False** at + runtime (expected). + * Delegation via ``__getattr__`` works for every member of the real context. + * No reliance on internal initialization side-effects of the durable SDK. + """ def __init__( self, @@ -38,14 +68,6 @@ def call_activity_with_retry( self._task_tracker.record_activity_call() return task - def set_custom_status(self, status: str): - """Set custom status for the orchestration.""" - self._context.set_custom_status(status) - - def wait_for_external_event(self, event_name: str): - """Wait for an external event in the orchestration.""" - return self._context.wait_for_external_event(event_name) - def activity_as_tool( self, activity_func: Callable, @@ -95,3 +117,14 @@ async def run_activity(ctx: RunContextWrapper[Any], input: str) -> Any: on_invoke_tool=run_activity, strict_json_schema=True, ) + + def __getattr__(self, name): + """Delegate missing attributes to the underlying DurableOrchestrationContext.""" + try: + return getattr(self._context, name) + except AttributeError: + raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'") + + def __dir__(self): + """Improve introspection and tab-completion by including delegated attributes.""" + return sorted(set(dir(type(self)) + list(self.__dict__) + dir(self._context))) diff --git a/tests/openai_agents/test_context.py b/tests/openai_agents/test_context.py new file mode 100644 index 00000000..0d4bbc8e --- /dev/null +++ b/tests/openai_agents/test_context.py @@ -0,0 +1,292 @@ +import pytest +from unittest.mock import Mock, patch + +from azure.durable_functions.openai_agents.context import DurableAIAgentContext +from azure.durable_functions.openai_agents.task_tracker import TaskTracker +from azure.durable_functions.models.DurableOrchestrationContext import DurableOrchestrationContext +from azure.durable_functions.models.RetryOptions import RetryOptions + +from agents.tool import FunctionTool + + +class TestDurableAIAgentContext: + """Test suite for DurableAIAgentContext class.""" + + def _create_mock_orchestration_context(self): + """Create a mock DurableOrchestrationContext for testing.""" + orchestration_context = Mock(spec=DurableOrchestrationContext) + orchestration_context.call_activity = Mock(return_value="mock_task") + orchestration_context.call_activity_with_retry = Mock(return_value="mock_task_with_retry") + orchestration_context.instance_id = "test_instance_id" + orchestration_context.current_utc_datetime = "2023-01-01T00:00:00Z" + orchestration_context.is_replaying = False + return orchestration_context + + def _create_mock_task_tracker(self): + """Create a mock TaskTracker for testing.""" + task_tracker = Mock(spec=TaskTracker) + task_tracker.record_activity_call = Mock() + task_tracker.get_activity_call_result = Mock(return_value="activity_result") + task_tracker.get_activity_call_result_with_retry = Mock(return_value="retry_activity_result") + return task_tracker + + def test_init_creates_context_successfully(self): + """Test that __init__ creates a DurableAIAgentContext successfully.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + retry_options = RetryOptions(1000, 3) + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, retry_options) + + assert isinstance(ai_context, DurableAIAgentContext) + assert not isinstance(ai_context, DurableOrchestrationContext) + + def test_call_activity_delegates_and_records(self): + """Test that call_activity delegates to context and records activity call.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + result = ai_context.call_activity("test_activity", "test_input") + + orchestration_context.call_activity.assert_called_once_with("test_activity", "test_input") + task_tracker.record_activity_call.assert_called_once() + assert result == "mock_task" + + def test_call_activity_with_retry_delegates_and_records(self): + """Test that call_activity_with_retry delegates to context and records activity call.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + retry_options = RetryOptions(1000, 3) + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + result = ai_context.call_activity_with_retry("test_activity", retry_options, "test_input") + + orchestration_context.call_activity_with_retry.assert_called_once_with( + "test_activity", retry_options, "test_input" + ) + task_tracker.record_activity_call.assert_called_once() + assert result == "mock_task_with_retry" + + @patch('azure.durable_functions.openai_agents.context.function_schema') + @patch('azure.durable_functions.openai_agents.context.FunctionTool') + def test_activity_as_tool_creates_function_tool(self, mock_function_tool, mock_function_schema): + """Test that activity_as_tool creates a FunctionTool with correct parameters.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + # Mock the activity function + mock_activity_func = Mock() + mock_activity_func._function._name = "test_activity" + mock_activity_func._function._func = lambda x: x + + # Mock the schema + mock_schema = Mock() + mock_schema.name = "test_activity" + mock_schema.description = "Test activity description" + mock_schema.params_json_schema = {"type": "object"} + mock_function_schema.return_value = mock_schema + + # Mock FunctionTool + mock_tool = Mock(spec=FunctionTool) + mock_function_tool.return_value = mock_tool + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + retry_options = RetryOptions(1000, 3) + + result = ai_context.activity_as_tool( + mock_activity_func, + description="Custom description", + retry_options=retry_options + ) + + # Verify function_schema was called correctly + mock_function_schema.assert_called_once_with( + func=mock_activity_func._function._func, + name_override="test_activity", + docstring_style=None, + description_override="Custom description", + use_docstring_info=True, + strict_json_schema=True, + ) + + # Verify FunctionTool was created correctly + mock_function_tool.assert_called_once() + call_args = mock_function_tool.call_args + assert call_args[1]['name'] == "test_activity" + assert call_args[1]['description'] == "Test activity description" + assert call_args[1]['params_json_schema'] == {"type": "object"} + assert call_args[1]['strict_json_schema'] is True + assert callable(call_args[1]['on_invoke_tool']) + + assert result is mock_tool + + @patch('azure.durable_functions.openai_agents.context.function_schema') + @patch('azure.durable_functions.openai_agents.context.FunctionTool') + def test_activity_as_tool_with_default_retry_options(self, mock_function_tool, mock_function_schema): + """Test that activity_as_tool uses default retry options when none provided.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + mock_activity_func = Mock() + mock_activity_func._function._name = "test_activity" + mock_activity_func._function._func = lambda x: x + + mock_schema = Mock() + mock_schema.name = "test_activity" + mock_schema.description = "Test description" + mock_schema.params_json_schema = {"type": "object"} + mock_function_schema.return_value = mock_schema + + mock_tool = Mock(spec=FunctionTool) + mock_function_tool.return_value = mock_tool + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + # Call with default retry options + result = ai_context.activity_as_tool(mock_activity_func) + + # Should still create the tool successfully + assert result is mock_tool + mock_function_tool.assert_called_once() + + @patch('azure.durable_functions.openai_agents.context.function_schema') + @patch('azure.durable_functions.openai_agents.context.FunctionTool') + def test_activity_as_tool_run_activity_with_retry(self, mock_function_tool, mock_function_schema): + """Test that the run_activity function calls task tracker with retry options.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + mock_activity_func = Mock() + mock_activity_func._function._name = "test_activity" + mock_activity_func._function._func = lambda x: x + + mock_schema = Mock() + mock_schema.name = "test_activity" + mock_schema.description = "" + mock_schema.params_json_schema = {"type": "object"} + mock_function_schema.return_value = mock_schema + + mock_tool = Mock(spec=FunctionTool) + mock_function_tool.return_value = mock_tool + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + retry_options = RetryOptions(1000, 3) + + ai_context.activity_as_tool(mock_activity_func, retry_options=retry_options) + + # Get the run_activity function that was passed to FunctionTool + call_args = mock_function_tool.call_args + run_activity = call_args[1]['on_invoke_tool'] + + # Create a mock context wrapper + mock_ctx = Mock() + + # Call the run_activity function + import asyncio + result = asyncio.run(run_activity(mock_ctx, "test_input")) + + # Verify the task tracker was called with retry options + task_tracker.get_activity_call_result_with_retry.assert_called_once_with( + "test_activity", retry_options, "test_input" + ) + assert result == "retry_activity_result" + + @patch('azure.durable_functions.openai_agents.context.function_schema') + @patch('azure.durable_functions.openai_agents.context.FunctionTool') + def test_activity_as_tool_run_activity_without_retry(self, mock_function_tool, mock_function_schema): + """Test that the run_activity function calls task tracker without retry when retry_options is None.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + mock_activity_func = Mock() + mock_activity_func._function._name = "test_activity" + mock_activity_func._function._func = lambda x: x + + mock_schema = Mock() + mock_schema.name = "test_activity" + mock_schema.description = "" + mock_schema.params_json_schema = {"type": "object"} + mock_function_schema.return_value = mock_schema + + mock_tool = Mock(spec=FunctionTool) + mock_function_tool.return_value = mock_tool + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + ai_context.activity_as_tool(mock_activity_func, retry_options=None) + + # Get the run_activity function that was passed to FunctionTool + call_args = mock_function_tool.call_args + run_activity = call_args[1]['on_invoke_tool'] + + # Create a mock context wrapper + mock_ctx = Mock() + + # Call the run_activity function + import asyncio + result = asyncio.run(run_activity(mock_ctx, "test_input")) + + # Verify the task tracker was called without retry options + task_tracker.get_activity_call_result.assert_called_once_with( + "test_activity", "test_input" + ) + assert result == "activity_result" + + def test_context_delegation_methods_work(self): + """Test that common context methods work through delegation.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + # Add some mock methods to the orchestration context + orchestration_context.wait_for_external_event = Mock(return_value="external_event_task") + orchestration_context.create_timer = Mock(return_value="timer_task") + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + # These should work through delegation + result1 = ai_context.wait_for_external_event("test_event") + result2 = ai_context.create_timer("2023-01-01T00:00:00Z") + + assert result1 == "external_event_task" + assert result2 == "timer_task" + orchestration_context.wait_for_external_event.assert_called_once_with("test_event") + orchestration_context.create_timer.assert_called_once_with("2023-01-01T00:00:00Z") + + def test_getattr_delegates_to_context(self): + """Test that __getattr__ delegates attribute access to the underlying context.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + # Test delegation of various attributes + assert ai_context.instance_id == "test_instance_id" + assert ai_context.current_utc_datetime == "2023-01-01T00:00:00Z" + assert ai_context.is_replaying is False + + def test_getattr_raises_attribute_error_for_nonexistent_attributes(self): + """Test that __getattr__ raises AttributeError for non-existent attributes.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + with pytest.raises(AttributeError, match="'DurableAIAgentContext' object has no attribute 'nonexistent_attr'"): + _ = ai_context.nonexistent_attr + + def test_dir_includes_delegated_attributes(self): + """Test that __dir__ includes attributes from the underlying context.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + dir_result = dir(ai_context) + + # Should include delegated attributes from the underlying context + assert 'instance_id' in dir_result + assert 'current_utc_datetime' in dir_result + assert 'is_replaying' in dir_result + # Should also include public methods + assert 'call_activity' in dir_result + assert 'activity_as_tool' in dir_result From b0d1ea30214385416dda2a35b8d9e03f9f92d991 Mon Sep 17 00:00:00 2001 From: Anatoli Beliaev Date: Fri, 12 Sep 2025 23:23:47 -0700 Subject: [PATCH 2/5] Fix linter issue --- azure/durable_functions/openai_agents/context.py | 1 + 1 file changed, 1 insertion(+) diff --git a/azure/durable_functions/openai_agents/context.py b/azure/durable_functions/openai_agents/context.py index 72657344..1957515a 100644 --- a/azure/durable_functions/openai_agents/context.py +++ b/azure/durable_functions/openai_agents/context.py @@ -22,6 +22,7 @@ class _BaseDurableContext: # lightweight runtime stub attribute lookups for its members are delegated via __getattr__ to the wrapped ``_context`` instance. """ + __slots__ = () From a3c10c75f79f9fdfd7e286d3962b92d89e766d9a Mon Sep 17 00:00:00 2001 From: Anatoli Beliaev Date: Tue, 16 Sep 2025 10:24:23 -0700 Subject: [PATCH 3/5] Rename `activity_as_tool` to `create_activity_tool` in tests --- tests/openai_agents/test_context.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/tests/openai_agents/test_context.py b/tests/openai_agents/test_context.py index 0d4bbc8e..1e74527d 100644 --- a/tests/openai_agents/test_context.py +++ b/tests/openai_agents/test_context.py @@ -71,7 +71,7 @@ def test_call_activity_with_retry_delegates_and_records(self): @patch('azure.durable_functions.openai_agents.context.function_schema') @patch('azure.durable_functions.openai_agents.context.FunctionTool') def test_activity_as_tool_creates_function_tool(self, mock_function_tool, mock_function_schema): - """Test that activity_as_tool creates a FunctionTool with correct parameters.""" + """Test that create_activity_tool creates a FunctionTool with correct parameters.""" orchestration_context = self._create_mock_orchestration_context() task_tracker = self._create_mock_task_tracker() @@ -94,7 +94,7 @@ def test_activity_as_tool_creates_function_tool(self, mock_function_tool, mock_f ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) retry_options = RetryOptions(1000, 3) - result = ai_context.activity_as_tool( + result = ai_context.create_activity_tool( mock_activity_func, description="Custom description", retry_options=retry_options @@ -124,7 +124,7 @@ def test_activity_as_tool_creates_function_tool(self, mock_function_tool, mock_f @patch('azure.durable_functions.openai_agents.context.function_schema') @patch('azure.durable_functions.openai_agents.context.FunctionTool') def test_activity_as_tool_with_default_retry_options(self, mock_function_tool, mock_function_schema): - """Test that activity_as_tool uses default retry options when none provided.""" + """Test that create_activity_tool uses default retry options when none provided.""" orchestration_context = self._create_mock_orchestration_context() task_tracker = self._create_mock_task_tracker() @@ -144,7 +144,7 @@ def test_activity_as_tool_with_default_retry_options(self, mock_function_tool, m ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) # Call with default retry options - result = ai_context.activity_as_tool(mock_activity_func) + result = ai_context.create_activity_tool(mock_activity_func) # Should still create the tool successfully assert result is mock_tool @@ -173,7 +173,7 @@ def test_activity_as_tool_run_activity_with_retry(self, mock_function_tool, mock ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) retry_options = RetryOptions(1000, 3) - ai_context.activity_as_tool(mock_activity_func, retry_options=retry_options) + ai_context.create_activity_tool(mock_activity_func, retry_options=retry_options) # Get the run_activity function that was passed to FunctionTool call_args = mock_function_tool.call_args @@ -214,7 +214,7 @@ def test_activity_as_tool_run_activity_without_retry(self, mock_function_tool, m ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) - ai_context.activity_as_tool(mock_activity_func, retry_options=None) + ai_context.create_activity_tool(mock_activity_func, retry_options=None) # Get the run_activity function that was passed to FunctionTool call_args = mock_function_tool.call_args @@ -289,4 +289,4 @@ def test_dir_includes_delegated_attributes(self): assert 'is_replaying' in dir_result # Should also include public methods assert 'call_activity' in dir_result - assert 'activity_as_tool' in dir_result + assert 'create_activity_tool' in dir_result From bea39adf8d993958a1869e2d2323398a34f34150 Mon Sep 17 00:00:00 2001 From: Anatoli Beliaev Date: Tue, 16 Sep 2025 10:53:49 -0700 Subject: [PATCH 4/5] Adjust create_activity_tool test expectations --- tests/openai_agents/test_context.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/openai_agents/test_context.py b/tests/openai_agents/test_context.py index 1e74527d..76d9cb36 100644 --- a/tests/openai_agents/test_context.py +++ b/tests/openai_agents/test_context.py @@ -103,7 +103,6 @@ def test_activity_as_tool_creates_function_tool(self, mock_function_tool, mock_f # Verify function_schema was called correctly mock_function_schema.assert_called_once_with( func=mock_activity_func._function._func, - name_override="test_activity", docstring_style=None, description_override="Custom description", use_docstring_info=True, @@ -159,6 +158,7 @@ def test_activity_as_tool_run_activity_with_retry(self, mock_function_tool, mock mock_activity_func = Mock() mock_activity_func._function._name = "test_activity" + mock_activity_func._function._trigger = None mock_activity_func._function._func = lambda x: x mock_schema = Mock() @@ -201,6 +201,7 @@ def test_activity_as_tool_run_activity_without_retry(self, mock_function_tool, m mock_activity_func = Mock() mock_activity_func._function._name = "test_activity" + mock_activity_func._function._trigger = None mock_activity_func._function._func = lambda x: x mock_schema = Mock() From 93cd42057385c467b39e61c466e624d7815ea9a4 Mon Sep 17 00:00:00 2001 From: Anatoli Beliaev Date: Tue, 16 Sep 2025 11:02:55 -0700 Subject: [PATCH 5/5] Add test_activity_as_tool_extracts_activity_name_from_trigger test --- tests/openai_agents/test_context.py | 42 +++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/openai_agents/test_context.py b/tests/openai_agents/test_context.py index 76d9cb36..cb887005 100644 --- a/tests/openai_agents/test_context.py +++ b/tests/openai_agents/test_context.py @@ -234,6 +234,48 @@ def test_activity_as_tool_run_activity_without_retry(self, mock_function_tool, m ) assert result == "activity_result" + @patch('azure.durable_functions.openai_agents.context.function_schema') + @patch('azure.durable_functions.openai_agents.context.FunctionTool') + def test_activity_as_tool_extracts_activity_name_from_trigger(self, mock_function_tool, mock_function_schema): + """Test that the run_activity function calls task tracker with the activity name specified in the trigger.""" + orchestration_context = self._create_mock_orchestration_context() + task_tracker = self._create_mock_task_tracker() + + mock_activity_func = Mock() + mock_activity_func._function._name = "test_activity" + mock_activity_func._function._trigger.activity = "activity_name_from_trigger" + mock_activity_func._function._func = lambda x: x + + mock_schema = Mock() + mock_schema.name = "test_activity" + mock_schema.description = "" + mock_schema.params_json_schema = {"type": "object"} + mock_function_schema.return_value = mock_schema + + mock_tool = Mock(spec=FunctionTool) + mock_function_tool.return_value = mock_tool + + ai_context = DurableAIAgentContext(orchestration_context, task_tracker, None) + + ai_context.create_activity_tool(mock_activity_func, retry_options=None) + + # Get the run_activity function that was passed to FunctionTool + call_args = mock_function_tool.call_args + run_activity = call_args[1]['on_invoke_tool'] + + # Create a mock context wrapper + mock_ctx = Mock() + + # Call the run_activity function + import asyncio + result = asyncio.run(run_activity(mock_ctx, "test_input")) + + # Verify the task tracker was called without retry options + task_tracker.get_activity_call_result.assert_called_once_with( + "activity_name_from_trigger", "test_input" + ) + assert result == "activity_result" + def test_context_delegation_methods_work(self): """Test that common context methods work through delegation.""" orchestration_context = self._create_mock_orchestration_context()