* Python: Add allowed_checkpoint_types support to CosmosCheckpointStorage (#5200) Add allowed_checkpoint_types parameter to CosmosCheckpointStorage for parity with FileCheckpointStorage. This ensures both providers use the same restricted pickle deserialization by default. Changes: - Accept allowed_checkpoint_types kwarg in __init__, stored as frozenset - Convert _document_to_checkpoint from @staticmethod to instance method - Forward allowed_types to decode_checkpoint_value on all load paths - Update class docstring to describe the new parameter - Add tests covering built-in safe types, app type opt-in/blocking, and all load paths (load, list_checkpoints, get_latest) - Add changelog entry noting the breaking behavior change BREAKING CHANGE: CosmosCheckpointStorage now uses restricted pickle deserialization by default. Checkpoints containing application-defined types will require passing those types via allowed_checkpoint_types. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> * Python: Add `allowed_checkpoint_types` support to `CosmosCheckpointStorage` for parity with `FileCheckpointStorage` Fixes #5200 * Address PR review: add pickle security warning and fix docstring examples - Reintroduce explicit security warning about pickle deserialization risks - Convert Example:: block to .. code-block:: python with imports for consistency with other docstring examples - Note: PR title should be updated to include [BREAKING] prefix per changelog convention (comment #3, requires GitHub UI change) Fixes #5200 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <copilot@github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Get Started with Microsoft Agent Framework Azure Cosmos DB
Please install this package via pip:
pip install agent-framework-azure-cosmos --pre
Azure Cosmos DB History Provider
The Azure Cosmos DB integration provides CosmosHistoryProvider for persistent conversation history storage.
Basic Usage Example
from azure.identity.aio import DefaultAzureCredential
from agent_framework_azure_cosmos import CosmosHistoryProvider
provider = CosmosHistoryProvider(
endpoint="https://<account>.documents.azure.com:443/",
credential=DefaultAzureCredential(),
database_name="agent-framework",
container_name="chat-history",
)
Credentials follow the same pattern used by other Azure connectors in the repository:
- Pass a credential object (for example
DefaultAzureCredential) - Or pass a key string directly
- Or set
AZURE_COSMOS_KEYin the environment
Container naming behavior:
- Container name is configured on the provider (
container_nameorAZURE_COSMOS_CONTAINER_NAME) session_idis used as the Cosmos partition key for reads/writes
See samples/02-agents/conversations/cosmos_history_provider.py for a runnable example.
Cosmos DB Workflow Checkpoint Storage
CosmosCheckpointStorage implements the CheckpointStorage protocol, enabling
durable workflow checkpointing backed by Azure Cosmos DB NoSQL. Workflows can be
paused and resumed across process restarts by persisting checkpoint state in Cosmos DB.
Basic Usage
Managed Identity / RBAC (recommended for production)
from azure.identity.aio import DefaultAzureCredential
from agent_framework import WorkflowBuilder
from agent_framework_azure_cosmos import CosmosCheckpointStorage
checkpoint_storage = CosmosCheckpointStorage(
endpoint="https://<account>.documents.azure.com:443/",
credential=DefaultAzureCredential(),
database_name="agent-framework",
container_name="workflow-checkpoints",
)
Account Key
from agent_framework_azure_cosmos import CosmosCheckpointStorage
checkpoint_storage = CosmosCheckpointStorage(
endpoint="https://<account>.documents.azure.com:443/",
credential="<your-account-key>",
database_name="agent-framework",
container_name="workflow-checkpoints",
)
Then use with a workflow
from agent_framework import WorkflowBuilder
# Build a workflow with checkpointing enabled
workflow = WorkflowBuilder(
start_executor=start,
checkpoint_storage=checkpoint_storage,
).build()
# Run the workflow — checkpoints are automatically saved after each superstep
result = await workflow.run(message="input data")
# Resume from a checkpoint
latest = await checkpoint_storage.get_latest(workflow_name=workflow.name)
if latest:
resumed = await workflow.run(checkpoint_id=latest.checkpoint_id)
Authentication Options
CosmosCheckpointStorage supports the same authentication modes as CosmosHistoryProvider:
- Managed identity / RBAC (recommended): Pass
DefaultAzureCredential(),ManagedIdentityCredential(), or any AzureTokenCredential - Account key: Pass a key string via
credentialparameter - Environment variables: Set
AZURE_COSMOS_ENDPOINT,AZURE_COSMOS_DATABASE_NAME,AZURE_COSMOS_CONTAINER_NAME, andAZURE_COSMOS_KEY(key not required when using Azure credentials) - Pre-created client: Pass an existing
CosmosClientorContainerProxy
Database and Container Setup
The database and container are created automatically on first use (via
create_database_if_not_exists and create_container_if_not_exists). The container
uses /workflow_name as the partition key. You can also pre-create them in the Azure
portal with this partition key configuration.
Environment Variables
| Variable | Description |
|---|---|
AZURE_COSMOS_ENDPOINT |
Cosmos DB account endpoint |
AZURE_COSMOS_DATABASE_NAME |
Database name |
AZURE_COSMOS_CONTAINER_NAME |
Container name |
AZURE_COSMOS_KEY |
Account key (optional if using Azure credentials) |
See samples/03-workflows/checkpoint/cosmos_workflow_checkpointing.py for a standalone example,
or samples/03-workflows/checkpoint/cosmos_workflow_checkpointing_foundry.py for an end-to-end
example with Azure AI Foundry agents.