[Hosted Agents] Implement managed checkpoints feature for Langgraph

sdk/agentserver/azure-ai-agentserver-agentframework/azure/ai/agentserver/agentframework/__init__.py

@@ -49,8 +49,6 @@ def from_agent_framework(
         credentials: Optional[Union[AsyncTokenCredential, TokenCredential]] = None,
         thread_repository: Optional[AgentThreadRepository] = None,
         checkpoint_repository: Optional[CheckpointRepository] = None,
-        managed_checkpoints: bool = False,
-        project_endpoint: Optional[str] = None,
     ) -> "AgentFrameworkWorkflowAdapter":
     """
     Create an Agent Framework Workflow Adapter.
@@ -68,13 +66,9 @@ def from_agent_framework(
     :param thread_repository: Optional thread repository for agent thread management.
     :type thread_repository: Optional[AgentThreadRepository]
     :param checkpoint_repository: Optional checkpoint repository for workflow checkpointing.
+        Use ``InMemoryCheckpointRepository``, ``FileCheckpointRepository``, or
+        ``FoundryCheckpointRepository`` for Azure AI Foundry managed storage.
     :type checkpoint_repository: Optional[CheckpointRepository]
-    :param managed_checkpoints: If True, use Azure AI Foundry managed checkpoint storage.
-    :type managed_checkpoints: bool
-    :param project_endpoint: The Azure AI Foundry project endpoint. If not provided,
-        will be read from AZURE_AI_PROJECT_ENDPOINT environment variable.
-        Example: "https://<resource>.services.ai.azure.com/api/projects/<project-id>"
-    :type project_endpoint: Optional[str]
     :return: An instance of AgentFrameworkWorkflowAdapter.
     :rtype: AgentFrameworkWorkflowAdapter
     """
@@ -86,8 +80,6 @@ def from_agent_framework(
     credentials: Optional[Union[AsyncTokenCredential, TokenCredential]] = None,
     thread_repository: Optional[AgentThreadRepository] = None,
     checkpoint_repository: Optional[CheckpointRepository] = None,
-    managed_checkpoints: bool = False,
-    project_endpoint: Optional[str] = None,
 ) -> "AgentFrameworkAgent":
     """
     Create an Agent Framework Adapter from either an AgentProtocol/BaseAgent or a
@@ -101,19 +93,13 @@ def from_agent_framework(
     :param thread_repository: Optional thread repository for agent thread management.
     :type thread_repository: Optional[AgentThreadRepository]
     :param checkpoint_repository: Optional checkpoint repository for workflow checkpointing.
+        Use ``InMemoryCheckpointRepository``, ``FileCheckpointRepository``, or
+        ``FoundryCheckpointRepository`` for Azure AI Foundry managed storage.
     :type checkpoint_repository: Optional[CheckpointRepository]
-    :param managed_checkpoints: If True, use Azure AI Foundry managed checkpoint storage.
-    :type managed_checkpoints: bool
-    :param project_endpoint: The Azure AI Foundry project endpoint. If not provided,
-        will be read from AZURE_AI_PROJECT_ENDPOINT environment variable.
-        Example: "https://<resource>.services.ai.azure.com/api/projects/<project-id>"
-    :type project_endpoint: Optional[str]
     :return: An instance of AgentFrameworkAgent.
     :rtype: AgentFrameworkAgent
     :raises TypeError: If neither or both of agent and workflow are provided, or if
                        the provided types are incorrect.
-    :raises ValueError: If managed_checkpoints=True but required parameters are missing,
-                       or if both managed_checkpoints=True and checkpoint_repository are provided.
     """
 
     if isinstance(agent_or_workflow, WorkflowBuilder):
@@ -122,17 +108,13 @@ def from_agent_framework(
             credentials=credentials,
             thread_repository=thread_repository,
             checkpoint_repository=checkpoint_repository,
-            managed_checkpoints=managed_checkpoints,
-            project_endpoint=project_endpoint,
         )
     if isinstance(agent_or_workflow, Callable):  # type: ignore
         return AgentFrameworkWorkflowAdapter(
             workflow_factory=agent_or_workflow,
             credentials=credentials,
             thread_repository=thread_repository,
             checkpoint_repository=checkpoint_repository,
-            managed_checkpoints=managed_checkpoints,
-            project_endpoint=project_endpoint,
         )
     # raise TypeError("workflow must be a WorkflowBuilder or callable returning a Workflow")
 

sdk/agentserver/azure-ai-agentserver-agentframework/azure/ai/agentserver/agentframework/_workflow_agent_adapter.py

@@ -16,7 +16,7 @@
 from azure.core.credentials_async import AsyncTokenCredential
 
 from azure.ai.agentserver.core import AgentRunContext
-from azure.ai.agentserver.core.logger import get_logger, get_project_endpoint
+from azure.ai.agentserver.core.logger import get_logger
 from azure.ai.agentserver.core.models import (
     Response as OpenAIResponse,
     ResponseStreamEvent,
@@ -28,7 +28,7 @@
 from .models.agent_framework_output_non_streaming_converter import (
     AgentFrameworkOutputNonStreamingConverter,
 )
-from .persistence import AgentThreadRepository, CheckpointRepository, FoundryCheckpointRepository
+from .persistence import AgentThreadRepository, CheckpointRepository
 
 logger = get_logger()
 
@@ -40,35 +40,9 @@ def __init__(
         credentials: Optional[Union[AsyncTokenCredential, TokenCredential]] = None,
         thread_repository: Optional[AgentThreadRepository] = None,
         checkpoint_repository: Optional[CheckpointRepository] = None,
-        managed_checkpoints: bool = False,
-        project_endpoint: Optional[str] = None,
     ) -> None:
         super().__init__(credentials, thread_repository)
         self._workflow_factory = workflow_factory
-
-        # Validate mutual exclusion of managed_checkpoints and checkpoint_repository
-        if managed_checkpoints and checkpoint_repository is not None:
-            raise ValueError(
-                "Cannot use both managed_checkpoints=True and checkpoint_repository. "
-                "Use managed_checkpoints=True for Azure AI Foundry managed storage, "
-                "or provide your own checkpoint_repository, but not both."
-            )
-
-        # Handle managed checkpoints
-        if managed_checkpoints:
-            resolved_endpoint = get_project_endpoint() or project_endpoint
-            if not resolved_endpoint:
-                raise ValueError(
-                    "project_endpoint is required when managed_checkpoints=True. "
-                    "Set AZURE_AI_PROJECT_ENDPOINT environment variable or pass project_endpoint parameter."
-                )
-            if not credentials:
-                raise ValueError("credentials are required when managed_checkpoints=True")
-            checkpoint_repository = FoundryCheckpointRepository(
-                project_endpoint=resolved_endpoint,
-                credential=credentials,
-            )
-
         self._checkpoint_repository = checkpoint_repository
 
     async def agent_run(  # pylint: disable=too-many-statements

sdk/agentserver/azure-ai-agentserver-agentframework/tests/unit_tests/test_from_agent_framework_managed.py

@@ -1,93 +1,47 @@
 # ---------------------------------------------------------
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
-"""Unit tests for from_agent_framework with managed checkpoints."""
+"""Unit tests for from_agent_framework with checkpoint repository."""
 
-import os
 import pytest
-from unittest.mock import Mock, AsyncMock, patch
+from unittest.mock import Mock
 
 from azure.core.credentials_async import AsyncTokenCredential
 
 
 @pytest.mark.unit
-def test_managed_checkpoints_requires_project_endpoint() -> None:
-    """Test that managed_checkpoints=True requires project_endpoint when env var not set."""
+def test_checkpoint_repository_is_optional() -> None:
+    """Test that checkpoint_repository is optional and defaults to None."""
     from azure.ai.agentserver.agentframework import from_agent_framework
     from agent_framework import WorkflowBuilder
 
     builder = WorkflowBuilder()
-    mock_credential = Mock(spec=AsyncTokenCredential)
 
-    # Ensure environment variable is not set
-    with patch.dict(os.environ, {}, clear=True):
-        with pytest.raises(ValueError) as exc_info:
-            from_agent_framework(
-                builder,
-                credentials=mock_credential,
-                managed_checkpoints=True,
-                project_endpoint=None,
-            )
+    # Should not raise
+    adapter = from_agent_framework(builder)
 
-    assert "project_endpoint" in str(exc_info.value)
+    assert adapter is not None
 
 
 @pytest.mark.unit
-def test_managed_checkpoints_requires_credentials() -> None:
-    """Test that managed_checkpoints=True requires credentials."""
+def test_foundry_checkpoint_repository_passed_directly() -> None:
+    """Test that FoundryCheckpointRepository can be passed via checkpoint_repository."""
     from azure.ai.agentserver.agentframework import from_agent_framework
+    from azure.ai.agentserver.agentframework.persistence import FoundryCheckpointRepository
     from agent_framework import WorkflowBuilder
 
     builder = WorkflowBuilder()
+    mock_credential = Mock(spec=AsyncTokenCredential)
 
-    with pytest.raises(ValueError) as exc_info:
-        from_agent_framework(
-            builder,
-            credentials=None,
-            managed_checkpoints=True,
-            project_endpoint="https://test.services.ai.azure.com/api/projects/test-project",
-        )
-
-    assert "credentials" in str(exc_info.value)
-
-
-@pytest.mark.unit
-def test_managed_checkpoints_false_does_not_require_parameters() -> None:
-    """Test that managed_checkpoints=False does not require project_endpoint."""
-    from azure.ai.agentserver.agentframework import from_agent_framework
-    from agent_framework import WorkflowBuilder
-
-    builder = WorkflowBuilder()
+    repo = FoundryCheckpointRepository(
+        project_endpoint="https://test.services.ai.azure.com/api/projects/test-project",
+        credential=mock_credential,
+    )
 
-    # Should not raise
     adapter = from_agent_framework(
         builder,
-        managed_checkpoints=False,
+        checkpoint_repository=repo,
     )
 
     assert adapter is not None
-
-
-@pytest.mark.unit
-def test_managed_checkpoints_and_checkpoint_repository_are_mutually_exclusive() -> None:
-    """Test that managed_checkpoints=True and checkpoint_repository cannot be used together."""
-    from azure.ai.agentserver.agentframework import from_agent_framework
-    from azure.ai.agentserver.agentframework.persistence import InMemoryCheckpointRepository
-    from agent_framework import WorkflowBuilder
-
-    builder = WorkflowBuilder()
-    mock_credential = Mock(spec=AsyncTokenCredential)
-    checkpoint_repo = InMemoryCheckpointRepository()
-
-    with pytest.raises(ValueError) as exc_info:
-        from_agent_framework(
-            builder,
-            credentials=mock_credential,
-            managed_checkpoints=True,
-            checkpoint_repository=checkpoint_repo,
-            project_endpoint="https://test.services.ai.azure.com/api/projects/test-project",
-        )
-
-    assert "Cannot use both" in str(exc_info.value)
-    assert "managed_checkpoints" in str(exc_info.value)
-    assert "checkpoint_repository" in str(exc_info.value)
+    assert adapter._checkpoint_repository is repo

sdk/agentserver/azure-ai-agentserver-langgraph/azure/ai/agentserver/langgraph/__init__.py

@@ -3,7 +3,7 @@
 # ---------------------------------------------------------
 __path__ = __import__("pkgutil").extend_path(__path__, __name__)
 
-from typing import Optional, TYPE_CHECKING
+from typing import Optional, Union, TYPE_CHECKING
 
 from azure.ai.agentserver.core.application import PackageMetadata, set_current_app
 
@@ -12,17 +12,30 @@
 from .langgraph import LangGraphAdapter
 
 if TYPE_CHECKING:  # pragma: no cover
+    from langgraph.graph.state import CompiledStateGraph
     from .models.response_api_converter import ResponseAPIConverter
     from azure.core.credentials_async import AsyncTokenCredential
+    from azure.core.credentials import TokenCredential
 
 
 def from_langgraph(
-    agent,
+    agent: "CompiledStateGraph",
     /,
-    credentials: Optional["AsyncTokenCredential"] = None,
-    converter: Optional["ResponseAPIConverter"] = None
+    credentials: Optional[Union["AsyncTokenCredential", "TokenCredential"]] = None,
+    converter: Optional["ResponseAPIConverter"] = None,
 ) -> "LangGraphAdapter":
-
+    """Create a LangGraph adapter for Azure AI Agent Server.
+
+    :param agent: The compiled LangGraph state graph. To use persistent checkpointing,
+        compile the graph with a checkpointer via ``builder.compile(checkpointer=saver)``.
+    :type agent: CompiledStateGraph
+    :param credentials: Azure credentials for authentication.
+    :type credentials: Optional[Union[AsyncTokenCredential, TokenCredential]]
+    :param converter: Custom response converter.
+    :type converter: Optional[ResponseAPIConverter]
+    :return: A LangGraphAdapter instance.
+    :rtype: LangGraphAdapter
+    """
     return LangGraphAdapter(agent, credentials=credentials, converter=converter)
 
 

sdk/agentserver/azure-ai-agentserver-langgraph/azure/ai/agentserver/langgraph/checkpointer/__init__.py

@@ -0,0 +1,8 @@
+# ---------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# ---------------------------------------------------------
+"""Checkpoint saver implementations for LangGraph with Azure AI Foundry."""
+
+from ._foundry_checkpoint_saver import FoundryCheckpointSaver
+
+__all__ = ["FoundryCheckpointSaver"]