chore: import upstream snapshot with attribution

python/packages/azure-cosmos/tests/test_cosmos_checkpoint_storage.py

@@ -0,0 +1,737 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import os
+import uuid
+from collections.abc import AsyncIterator
+from contextlib import suppress
+from dataclasses import dataclass
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from agent_framework._workflows._checkpoint import WorkflowCheckpoint
+from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
+from agent_framework.exceptions import SettingNotFoundError, WorkflowCheckpointException
+from azure.cosmos.aio import CosmosClient
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+import agent_framework_azure_cosmos._checkpoint_storage as checkpoint_storage_module
+from agent_framework_azure_cosmos._checkpoint_storage import CosmosCheckpointStorage
+
+skip_if_cosmos_integration_tests_disabled = pytest.mark.skipif(
+    any(
+        os.getenv(name, "") == ""
+        for name in (
+            "AZURE_COSMOS_ENDPOINT",
+            "AZURE_COSMOS_KEY",
+            "AZURE_COSMOS_DATABASE_NAME",
+            "AZURE_COSMOS_CONTAINER_NAME",
+        )
+    ),
+    reason=(
+        "AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_KEY, AZURE_COSMOS_DATABASE_NAME, and "
+        "AZURE_COSMOS_CONTAINER_NAME are required for Cosmos integration tests."
+    ),
+)
+
+
+def _to_async_iter(items: list[Any]) -> AsyncIterator[Any]:
+    async def _iterator() -> AsyncIterator[Any]:
+        for item in items:
+            yield item
+
+    return _iterator()
+
+
+def _make_checkpoint(
+    workflow_name: str = "test-workflow",
+    checkpoint_id: str | None = None,
+    previous_checkpoint_id: str | None = None,
+    timestamp: str | None = None,
+) -> WorkflowCheckpoint:
+    """Create a minimal WorkflowCheckpoint for testing."""
+    return WorkflowCheckpoint(
+        workflow_name=workflow_name,
+        graph_signature_hash="abc123",
+        checkpoint_id=checkpoint_id or str(uuid.uuid4()),
+        previous_checkpoint_id=previous_checkpoint_id,
+        timestamp=timestamp or "2025-01-01T00:00:00+00:00",
+        state={"counter": 42},
+        iteration_count=1,
+    )
+
+
+def _checkpoint_to_cosmos_document(checkpoint: WorkflowCheckpoint) -> dict[str, Any]:
+    """Simulate what a Cosmos DB document looks like after save."""
+    encoded = encode_checkpoint_value(checkpoint.to_dict())
+    doc: dict[str, Any] = {
+        "id": f"{checkpoint.workflow_name}_{checkpoint.checkpoint_id}",
+        "workflow_name": checkpoint.workflow_name,
+        **encoded,
+        # Cosmos system properties
+        "_rid": "abc",
+        "_self": "dbs/abc/colls/def/docs/ghi",
+        "_etag": '"00000000-0000-0000-0000-000000000000"',
+        "_attachments": "attachments/",
+        "_ts": 1700000000,
+    }
+    return doc
+
+
+@pytest.fixture
+def mock_container() -> MagicMock:
+    container = MagicMock()
+    container.query_items = MagicMock(return_value=_to_async_iter([]))
+    container.upsert_item = AsyncMock(return_value={})
+    container.delete_item = AsyncMock(return_value={})
+    return container
+
+
+@pytest.fixture
+def mock_cosmos_client(mock_container: MagicMock) -> MagicMock:
+    database_client = MagicMock()
+    database_client.create_container_if_not_exists = AsyncMock(return_value=mock_container)
+
+    client = MagicMock()
+    client.create_database_if_not_exists = AsyncMock(return_value=database_client)
+    client.close = AsyncMock()
+    return client
+
+
+# --- Tests for initialization ---
+
+
+async def test_init_uses_provided_container_client(mock_container: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    assert storage.database_name == ""
+    assert storage.container_name == ""
+
+
+async def test_init_uses_provided_cosmos_client(mock_cosmos_client: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(
+        cosmos_client=mock_cosmos_client,
+        database_name="db1",
+        container_name="checkpoints",
+    )
+    assert storage.database_name == "db1"
+    assert storage.container_name == "checkpoints"
+
+
+async def test_init_missing_required_settings_raises(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.delenv("AZURE_COSMOS_ENDPOINT", raising=False)
+    monkeypatch.delenv("AZURE_COSMOS_DATABASE_NAME", raising=False)
+    monkeypatch.delenv("AZURE_COSMOS_CONTAINER_NAME", raising=False)
+    monkeypatch.delenv("AZURE_COSMOS_KEY", raising=False)
+
+    with pytest.raises(SettingNotFoundError, match="database_name"):
+        CosmosCheckpointStorage()
+
+
+async def test_init_constructs_client_with_credential(
+    monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock
+) -> None:
+    """Uses key-based auth when a key string is provided, otherwise falls back to Azure credential (RBAC)."""
+    mock_factory = MagicMock(return_value=mock_cosmos_client)
+    monkeypatch.setattr(checkpoint_storage_module, "CosmosClient", mock_factory)
+    monkeypatch.delenv("AZURE_COSMOS_KEY", raising=False)
+
+    # Simulate real-world pattern: use key if available, else RBAC credential
+    cosmos_key = os.getenv("AZURE_COSMOS_KEY")
+    credential: Any = cosmos_key if cosmos_key else MagicMock()  # MagicMock simulates DefaultAzureCredential()
+
+    CosmosCheckpointStorage(
+        endpoint="https://account.documents.azure.com:443/",
+        credential=credential,
+        database_name="db1",
+        container_name="checkpoints",
+    )
+
+    mock_factory.assert_called_once()
+    kwargs = mock_factory.call_args.kwargs
+    assert kwargs["url"] == "https://account.documents.azure.com:443/"
+    assert kwargs["credential"] is credential
+
+
+async def test_init_creates_database_and_container(mock_cosmos_client: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(
+        cosmos_client=mock_cosmos_client,
+        database_name="db1",
+        container_name="custom-checkpoints",
+    )
+
+    await storage.list_checkpoint_ids(workflow_name="wf")
+
+    mock_cosmos_client.create_database_if_not_exists.assert_awaited_once_with(id="db1")
+    database_client = mock_cosmos_client.create_database_if_not_exists.return_value
+    assert database_client.create_container_if_not_exists.await_count == 1
+    kwargs = database_client.create_container_if_not_exists.await_args.kwargs
+    assert kwargs["id"] == "custom-checkpoints"
+
+
+# --- Tests for save ---
+
+
+async def test_save_upserts_document(mock_container: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    checkpoint = _make_checkpoint()
+
+    result = await storage.save(checkpoint)
+
+    assert result == checkpoint.checkpoint_id
+    mock_container.upsert_item.assert_awaited_once()
+    document = mock_container.upsert_item.await_args.kwargs["body"]
+    assert document["id"] == f"test-workflow_{checkpoint.checkpoint_id}"
+    assert document["workflow_name"] == "test-workflow"
+    assert document["graph_signature_hash"] == "abc123"
+    assert document["state"]["counter"] == 42
+
+
+async def test_save_returns_checkpoint_id(mock_container: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    checkpoint = _make_checkpoint(checkpoint_id="cp-123")
+
+    result = await storage.save(checkpoint)
+
+    assert result == "cp-123"
+
+
+# --- Tests for load ---
+
+
+async def test_load_returns_checkpoint(mock_container: MagicMock) -> None:
+    checkpoint = _make_checkpoint(checkpoint_id="cp-load")
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    loaded = await storage.load("cp-load")
+
+    assert loaded.checkpoint_id == "cp-load"
+    assert loaded.workflow_name == "test-workflow"
+    assert loaded.graph_signature_hash == "abc123"
+    assert loaded.state["counter"] == 42
+
+
+async def test_load_nonexistent_raises(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with pytest.raises(WorkflowCheckpointException, match="No checkpoint found"):
+        await storage.load("nonexistent-id")
+
+
+async def test_load_queries_without_partition_key(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    with suppress(WorkflowCheckpointException):
+        await storage.load("cp-id")
+
+    kwargs = mock_container.query_items.call_args.kwargs
+    assert "partition_key" not in kwargs
+
+
+async def test_load_multiple_workflows_same_checkpoint_id_raises(mock_container: MagicMock) -> None:
+    cp1 = _make_checkpoint(checkpoint_id="shared-id", workflow_name="workflow-a")
+    cp2 = _make_checkpoint(checkpoint_id="shared-id", workflow_name="workflow-b")
+    mock_container.query_items.return_value = _to_async_iter([
+        _checkpoint_to_cosmos_document(cp1),
+        _checkpoint_to_cosmos_document(cp2),
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with pytest.raises(WorkflowCheckpointException, match="Multiple checkpoints found"):
+        await storage.load("shared-id")
+
+
+# --- Tests for list_checkpoints ---
+
+
+async def test_list_checkpoints_returns_checkpoints_for_workflow(mock_container: MagicMock) -> None:
+    cp1 = _make_checkpoint(checkpoint_id="cp-1", timestamp="2025-01-01T00:00:00+00:00")
+    cp2 = _make_checkpoint(checkpoint_id="cp-2", timestamp="2025-01-02T00:00:00+00:00")
+    mock_container.query_items.return_value = _to_async_iter([
+        _checkpoint_to_cosmos_document(cp1),
+        _checkpoint_to_cosmos_document(cp2),
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    results = await storage.list_checkpoints(workflow_name="test-workflow")
+
+    assert len(results) == 2
+    assert results[0].checkpoint_id == "cp-1"
+    assert results[1].checkpoint_id == "cp-2"
+
+
+async def test_list_checkpoints_uses_partition_key(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    await storage.list_checkpoints(workflow_name="my-workflow")
+
+    kwargs = mock_container.query_items.call_args.kwargs
+    assert kwargs["partition_key"] == "my-workflow"
+
+
+async def test_list_checkpoints_empty_returns_empty(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    results = await storage.list_checkpoints(workflow_name="test-workflow")
+
+    assert results == []
+
+
+async def test_list_checkpoints_skips_malformed_documents(mock_container: MagicMock) -> None:
+    valid_cp = _make_checkpoint(checkpoint_id="cp-valid")
+    mock_container.query_items.return_value = _to_async_iter([
+        {"id": "bad_doc", "workflow_name": "test-workflow", "not_a_checkpoint": True},
+        _checkpoint_to_cosmos_document(valid_cp),
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    results = await storage.list_checkpoints(workflow_name="test-workflow")
+
+    assert len(results) == 1
+    assert results[0].checkpoint_id == "cp-valid"
+
+
+# --- Tests for delete ---
+
+
+async def test_delete_existing_returns_true(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([
+        {"id": "test-workflow_cp-del", "workflow_name": "test-workflow"},
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    result = await storage.delete("cp-del")
+
+    assert result is True
+    mock_container.delete_item.assert_awaited_once_with(
+        item="test-workflow_cp-del",
+        partition_key="test-workflow",
+    )
+
+
+async def test_delete_nonexistent_returns_false(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    result = await storage.delete("nonexistent")
+
+    assert result is False
+    mock_container.delete_item.assert_not_awaited()
+
+
+async def test_delete_cosmos_not_found_returns_false(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([
+        {"id": "test-workflow_cp-del", "workflow_name": "test-workflow"},
+    ])
+    mock_container.delete_item = AsyncMock(side_effect=CosmosResourceNotFoundError)
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    result = await storage.delete("cp-del")
+
+    assert result is False
+
+
+# --- Tests for get_latest ---
+
+
+async def test_get_latest_returns_latest_checkpoint(mock_container: MagicMock) -> None:
+    cp = _make_checkpoint(checkpoint_id="cp-latest", timestamp="2025-06-01T00:00:00+00:00")
+    mock_container.query_items.return_value = _to_async_iter([
+        _checkpoint_to_cosmos_document(cp),
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    result = await storage.get_latest(workflow_name="test-workflow")
+
+    assert result is not None
+    assert result.checkpoint_id == "cp-latest"
+
+
+async def test_get_latest_returns_none_when_empty(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    result = await storage.get_latest(workflow_name="test-workflow")
+
+    assert result is None
+
+
+async def test_get_latest_uses_order_by_desc_with_limit(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    await storage.get_latest(workflow_name="test-workflow")
+
+    kwargs = mock_container.query_items.call_args.kwargs
+    assert "ORDER BY c.timestamp DESC" in kwargs["query"]
+    assert "OFFSET 0 LIMIT 1" in kwargs["query"]
+
+
+# --- Tests for list_checkpoint_ids ---
+
+
+async def test_list_checkpoint_ids_returns_ids(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([
+        {"checkpoint_id": "cp-1"},
+        {"checkpoint_id": "cp-2"},
+    ])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    ids = await storage.list_checkpoint_ids(workflow_name="test-workflow")
+
+    assert ids == ["cp-1", "cp-2"]
+
+
+async def test_list_checkpoint_ids_empty_returns_empty(mock_container: MagicMock) -> None:
+    mock_container.query_items.return_value = _to_async_iter([])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    ids = await storage.list_checkpoint_ids(workflow_name="test-workflow")
+
+    assert ids == []
+
+
+# --- Tests for close and context manager ---
+
+
+async def test_close_closes_owned_client(monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock) -> None:
+    mock_factory = MagicMock(return_value=mock_cosmos_client)
+    monkeypatch.setattr(checkpoint_storage_module, "CosmosClient", mock_factory)
+
+    storage = CosmosCheckpointStorage(
+        endpoint="https://account.documents.azure.com:443/",
+        credential="key-123",
+        database_name="db1",
+        container_name="checkpoints",
+    )
+
+    await storage.close()
+
+    mock_cosmos_client.close.assert_awaited_once()
+
+
+async def test_close_does_not_close_external_client(mock_cosmos_client: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(
+        cosmos_client=mock_cosmos_client,
+        database_name="db1",
+        container_name="checkpoints",
+    )
+
+    await storage.close()
+
+    mock_cosmos_client.close.assert_not_awaited()
+
+
+async def test_context_manager_closes_owned_client(
+    monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock
+) -> None:
+    mock_factory = MagicMock(return_value=mock_cosmos_client)
+    monkeypatch.setattr(checkpoint_storage_module, "CosmosClient", mock_factory)
+
+    async with CosmosCheckpointStorage(
+        endpoint="https://account.documents.azure.com:443/",
+        credential="key-123",
+        database_name="db1",
+        container_name="checkpoints",
+    ) as storage:
+        assert storage is not None
+
+    mock_cosmos_client.close.assert_awaited_once()
+
+
+async def test_context_manager_preserves_original_exception(mock_container: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with (
+        patch.object(storage, "close", AsyncMock(side_effect=RuntimeError("close failed"))),
+        pytest.raises(ValueError, match="inner error"),
+    ):
+        async with storage:
+            raise ValueError("inner error")
+
+
+async def test_context_manager_reraises_close_error(mock_container: MagicMock) -> None:
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with (
+        patch.object(storage, "close", AsyncMock(side_effect=RuntimeError("close failed"))),
+        pytest.raises(RuntimeError, match="close failed"),
+    ):
+        async with storage:
+            pass  # no inner exception — close error should propagate
+
+
+# --- Tests for save/load round-trip ---
+
+
+async def test_round_trip_preserves_data(mock_container: MagicMock) -> None:
+    checkpoint = _make_checkpoint(
+        checkpoint_id="cp-roundtrip",
+        previous_checkpoint_id="cp-parent",
+    )
+    checkpoint.state = {"key": "value", "nested": {"a": 1}}
+    checkpoint.metadata = {"superstep": 3}
+    checkpoint.iteration_count = 5
+
+    saved_doc: dict[str, Any] = {}
+
+    async def capture_upsert(body: dict[str, Any]) -> dict[str, Any]:
+        saved_doc.update(body)
+        return body
+
+    mock_container.upsert_item = AsyncMock(side_effect=capture_upsert)
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    await storage.save(checkpoint)
+
+    returned_doc = {
+        **saved_doc,
+        "_rid": "abc",
+        "_self": "dbs/abc/colls/def/docs/ghi",
+        "_etag": '"etag"',
+        "_attachments": "attachments/",
+        "_ts": 1700000000,
+    }
+    mock_container.query_items.return_value = _to_async_iter([returned_doc])
+
+    loaded = await storage.load("cp-roundtrip")
+
+    assert loaded.checkpoint_id == checkpoint.checkpoint_id
+    assert loaded.workflow_name == checkpoint.workflow_name
+    assert loaded.graph_signature_hash == checkpoint.graph_signature_hash
+    assert loaded.previous_checkpoint_id == "cp-parent"
+    assert loaded.state == {"key": "value", "nested": {"a": 1}}
+    assert loaded.metadata == {"superstep": 3}
+    assert loaded.iteration_count == 5
+    assert loaded.version == "1.0"
+
+
+# --- Integration test ---
+
+
+@pytest.mark.integration
+@skip_if_cosmos_integration_tests_disabled
+async def test_cosmos_checkpoint_storage_roundtrip_with_emulator() -> None:
+    endpoint = os.getenv("AZURE_COSMOS_ENDPOINT", "")
+    key = os.getenv("AZURE_COSMOS_KEY", "")
+    database_prefix = os.getenv("AZURE_COSMOS_DATABASE_NAME", "")
+    container_prefix = os.getenv("AZURE_COSMOS_CONTAINER_NAME", "")
+    unique = uuid.uuid4().hex[:8]
+    database_name = f"{database_prefix}-cp-{unique}"
+    container_name = f"{container_prefix}-cp-{unique}"
+
+    async with CosmosClient(url=endpoint, credential=key) as cosmos_client:
+        await cosmos_client.create_database_if_not_exists(id=database_name)
+
+        storage = CosmosCheckpointStorage(
+            cosmos_client=cosmos_client,
+            database_name=database_name,
+            container_name=container_name,
+        )
+
+        try:
+            # Save two checkpoints for the same workflow
+            cp1 = _make_checkpoint(
+                checkpoint_id="cp-int-1",
+                workflow_name="integration-wf",
+                timestamp="2025-01-01T00:00:00+00:00",
+            )
+            cp2 = _make_checkpoint(
+                checkpoint_id="cp-int-2",
+                workflow_name="integration-wf",
+                previous_checkpoint_id="cp-int-1",
+                timestamp="2025-01-02T00:00:00+00:00",
+            )
+            cp2.state = {"step": 2}
+
+            await storage.save(cp1)
+            await storage.save(cp2)
+
+            # Load by ID
+            loaded = await storage.load("cp-int-1")
+            assert loaded.checkpoint_id == "cp-int-1"
+            assert loaded.workflow_name == "integration-wf"
+
+            # List all checkpoints for workflow
+            all_cps = await storage.list_checkpoints(workflow_name="integration-wf")
+            assert len(all_cps) == 2
+
+            # List checkpoint IDs
+            ids = await storage.list_checkpoint_ids(workflow_name="integration-wf")
+            assert "cp-int-1" in ids
+            assert "cp-int-2" in ids
+
+            # Get latest
+            latest = await storage.get_latest(workflow_name="integration-wf")
+            assert latest is not None
+            assert latest.checkpoint_id == "cp-int-2"
+            assert latest.state == {"step": 2}
+
+            # Delete
+            assert await storage.delete("cp-int-1") is True
+            assert await storage.delete("cp-int-1") is False
+
+            remaining = await storage.list_checkpoint_ids(workflow_name="integration-wf")
+            assert remaining == ["cp-int-2"]
+
+            # Cross-workflow isolation
+            other_cp = _make_checkpoint(
+                checkpoint_id="cp-other",
+                workflow_name="other-wf",
+            )
+            await storage.save(other_cp)
+            wf_cps = await storage.list_checkpoints(workflow_name="integration-wf")
+            assert len(wf_cps) == 1
+            assert wf_cps[0].checkpoint_id == "cp-int-2"
+
+        finally:
+            with suppress(Exception):
+                await cosmos_client.delete_database(database_name)
+
+
+# --- Tests for allowed_checkpoint_types ---
+
+
+@dataclass
+class _AppState:
+    """Application-defined state type used to test allowed_checkpoint_types."""
+
+    label: str
+    count: int
+
+
+_APP_STATE_TYPE_KEY = f"{_AppState.__module__}:{_AppState.__qualname__}"
+
+
+def _make_checkpoint_with_state(state: dict[str, Any]) -> WorkflowCheckpoint:
+    """Create a checkpoint with custom state for serialization tests."""
+    return WorkflowCheckpoint(
+        workflow_name="test-workflow",
+        graph_signature_hash="abc123",
+        timestamp="2025-01-01T00:00:00+00:00",
+        state=state,
+        iteration_count=1,
+    )
+
+
+async def test_init_accepts_allowed_checkpoint_types(mock_container: MagicMock) -> None:
+    """CosmosCheckpointStorage.__init__ accepts allowed_checkpoint_types."""
+    storage = CosmosCheckpointStorage(
+        container_client=mock_container,
+        allowed_checkpoint_types=["some.module:SomeType"],
+    )
+    assert storage is not None
+
+
+async def test_load_allows_builtin_safe_types(mock_container: MagicMock) -> None:
+    """Built-in safe types load without opt-in via allowed_checkpoint_types."""
+    from datetime import datetime, timezone
+
+    checkpoint = _make_checkpoint_with_state({
+        "ts": datetime(2025, 1, 1, tzinfo=timezone.utc),
+        "tags": {1, 2, 3},
+    })
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    loaded = await storage.load(checkpoint.checkpoint_id)
+
+    assert loaded.state["ts"] == datetime(2025, 1, 1, tzinfo=timezone.utc)
+    assert loaded.state["tags"] == {1, 2, 3}
+
+
+async def test_load_blocks_unlisted_app_type(mock_container: MagicMock) -> None:
+    """Application types are blocked when not listed in allowed_checkpoint_types."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        await storage.load(checkpoint.checkpoint_id)
+
+
+async def test_load_allows_listed_app_type(mock_container: MagicMock) -> None:
+    """Application types are allowed when listed in allowed_checkpoint_types."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="ok", count=7)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(
+        container_client=mock_container,
+        allowed_checkpoint_types=[_APP_STATE_TYPE_KEY],
+    )
+    loaded = await storage.load(checkpoint.checkpoint_id)
+
+    assert isinstance(loaded.state["data"], _AppState)
+    assert loaded.state["data"].label == "ok"
+    assert loaded.state["data"].count == 7
+
+
+async def test_list_checkpoints_blocks_unlisted_app_type(mock_container: MagicMock) -> None:
+    """list_checkpoints skips documents with unlisted application types."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+    results = await storage.list_checkpoints(workflow_name="test-workflow")
+
+    # The document is skipped (logged as warning) because the type is blocked
+    assert len(results) == 0
+
+
+async def test_list_checkpoints_allows_listed_app_type(mock_container: MagicMock) -> None:
+    """list_checkpoints decodes documents with listed application types."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="ok", count=3)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(
+        container_client=mock_container,
+        allowed_checkpoint_types=[_APP_STATE_TYPE_KEY],
+    )
+    results = await storage.list_checkpoints(workflow_name="test-workflow")
+
+    assert len(results) == 1
+    assert isinstance(results[0].state["data"], _AppState)
+
+
+async def test_get_latest_blocks_unlisted_app_type(mock_container: MagicMock) -> None:
+    """get_latest raises when the checkpoint contains an unlisted application type."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="x", count=1)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(container_client=mock_container)
+
+    with pytest.raises(WorkflowCheckpointException, match="deserialization blocked"):
+        await storage.get_latest(workflow_name="test-workflow")
+
+
+async def test_get_latest_allows_listed_app_type(mock_container: MagicMock) -> None:
+    """get_latest decodes checkpoints with listed application types."""
+    checkpoint = _make_checkpoint_with_state({"data": _AppState(label="latest", count=9)})
+    doc = _checkpoint_to_cosmos_document(checkpoint)
+    mock_container.query_items.return_value = _to_async_iter([doc])
+
+    storage = CosmosCheckpointStorage(
+        container_client=mock_container,
+        allowed_checkpoint_types=[_APP_STATE_TYPE_KEY],
+    )
+    result = await storage.get_latest(workflow_name="test-workflow")
+
+    assert result is not None
+    assert isinstance(result.state["data"], _AppState)
+    assert result.state["data"].label == "latest"

python/packages/azure-cosmos/tests/test_cosmos_history_provider.py

@@ -0,0 +1,417 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+import os
+import uuid
+from collections.abc import AsyncIterator
+from contextlib import suppress
+from typing import Any, cast
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from agent_framework import AgentResponse, Message
+from agent_framework._sessions import AgentSession, SessionContext
+from agent_framework.exceptions import SettingNotFoundError
+from azure.cosmos.aio import CosmosClient
+from azure.cosmos.exceptions import CosmosResourceNotFoundError
+
+import agent_framework_azure_cosmos._history_provider as history_provider_module
+from agent_framework_azure_cosmos._history_provider import CosmosHistoryProvider
+
+skip_if_cosmos_integration_tests_disabled = pytest.mark.skipif(
+    any(
+        os.getenv(name, "") == ""
+        for name in (
+            "AZURE_COSMOS_ENDPOINT",
+            "AZURE_COSMOS_KEY",
+            "AZURE_COSMOS_DATABASE_NAME",
+            "AZURE_COSMOS_CONTAINER_NAME",
+        )
+    ),
+    reason=(
+        "AZURE_COSMOS_ENDPOINT, AZURE_COSMOS_KEY, AZURE_COSMOS_DATABASE_NAME, and "
+        "AZURE_COSMOS_CONTAINER_NAME are required for Cosmos integration tests."
+    ),
+)
+
+
+def _to_async_iter(items: list[Any]) -> AsyncIterator[Any]:
+    async def _iterator() -> AsyncIterator[Any]:
+        for item in items:
+            yield item
+
+    return _iterator()
+
+
+@pytest.fixture
+def mock_container() -> MagicMock:
+    container = MagicMock()
+    container.query_items = MagicMock(return_value=_to_async_iter([]))
+    container.execute_item_batch = AsyncMock(return_value=[])
+    return container
+
+
+@pytest.fixture
+def mock_cosmos_client(mock_container: MagicMock) -> MagicMock:
+    database_client = MagicMock()
+    database_client.create_container_if_not_exists = AsyncMock(return_value=mock_container)
+
+    client = MagicMock()
+    client.get_database_client.return_value = database_client
+    client.close = AsyncMock()
+    return client
+
+
+class TestCosmosHistoryProviderInit:
+    def test_uses_provided_container_client(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        assert provider.source_id == "mem"
+        assert provider.load_messages is True
+        assert provider.store_outputs is True
+        assert provider.store_inputs is True
+        assert provider.database_name == ""
+        assert provider.container_name == ""
+
+    def test_uses_provided_cosmos_client(self, mock_cosmos_client: MagicMock) -> None:
+        provider = CosmosHistoryProvider(
+            source_id="mem",
+            cosmos_client=mock_cosmos_client,
+            database_name="db1",
+            container_name="history",
+        )
+
+        mock_cosmos_client.get_database_client.assert_called_once_with("db1")
+        assert provider.database_name == "db1"
+        assert provider.container_name == "history"
+
+    def test_missing_required_settings_raises(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.delenv("AZURE_COSMOS_ENDPOINT", raising=False)
+        monkeypatch.delenv("AZURE_COSMOS_DATABASE_NAME", raising=False)
+        monkeypatch.delenv("AZURE_COSMOS_CONTAINER_NAME", raising=False)
+        monkeypatch.delenv("AZURE_COSMOS_KEY", raising=False)
+
+        with pytest.raises(SettingNotFoundError, match="database_name"):
+            CosmosHistoryProvider()
+
+    def test_constructs_client_with_string_credential(
+        self, monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock
+    ) -> None:
+        mock_factory = MagicMock(return_value=mock_cosmos_client)
+        monkeypatch.setattr(history_provider_module, "CosmosClient", mock_factory)
+
+        CosmosHistoryProvider(
+            endpoint="https://account.documents.azure.com:443/",
+            credential="key-123",
+            database_name="db1",
+            container_name="history",
+        )
+
+        mock_factory.assert_called_once()
+        kwargs = mock_factory.call_args.kwargs
+        assert kwargs["url"] == "https://account.documents.azure.com:443/"
+        assert kwargs["credential"] == "key-123"
+
+
+class TestCosmosHistoryProviderContainerConfig:
+    async def test_provider_container_name_is_used(self, mock_cosmos_client: MagicMock) -> None:
+        provider = CosmosHistoryProvider(
+            source_id="mem",
+            cosmos_client=mock_cosmos_client,
+            database_name="db1",
+            container_name="custom-history",
+        )
+
+        await provider.get_messages("session-123")
+
+        database_client = mock_cosmos_client.get_database_client.return_value
+        assert database_client.create_container_if_not_exists.await_count == 1
+        kwargs = database_client.create_container_if_not_exists.await_args.kwargs
+        assert kwargs["id"] == "custom-history"
+
+
+class TestCosmosHistoryProviderGetMessages:
+    async def test_returns_deserialized_messages(self, mock_container: MagicMock) -> None:
+        msg1 = Message(role="user", contents=["Hello"])
+        msg2 = Message(role="assistant", contents=["Hi"])
+        mock_container.query_items.return_value = _to_async_iter([
+            {"message": msg1.to_dict()},
+            {"message": msg2.to_dict()},
+        ])
+
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        messages = await provider.get_messages("s1")
+
+        assert len(messages) == 2
+        assert messages[0].role == "user"
+        assert messages[0].text == "Hello"
+        assert messages[1].role == "assistant"
+        assert messages[1].text == "Hi"
+        query_kwargs = mock_container.query_items.call_args.kwargs
+        assert query_kwargs["partition_key"] == "s1"
+        assert query_kwargs["query"] == (
+            "SELECT c.message FROM c "
+            "WHERE c.session_id = @session_id AND c.source_id = @source_id "
+            "ORDER BY c.sort_key ASC"
+        )
+        assert query_kwargs["parameters"] == [
+            {"name": "@session_id", "value": "s1"},
+            {"name": "@source_id", "value": "mem"},
+        ]
+
+    async def test_empty_returns_empty(self, mock_container: MagicMock) -> None:
+        mock_container.query_items.return_value = _to_async_iter([])
+
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        messages = await provider.get_messages("s1")
+
+        assert messages == []
+
+    async def test_none_session_id_generates_guid_partition_key(
+        self, mock_container: MagicMock, caplog: pytest.LogCaptureFixture
+    ) -> None:
+        mock_container.query_items.return_value = _to_async_iter([])
+
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        with caplog.at_level("WARNING"):
+            await provider.get_messages(None)
+
+        query_kwargs = mock_container.query_items.call_args.kwargs
+        session_key = query_kwargs["partition_key"]
+        assert isinstance(session_key, str)
+        assert session_key != ""
+        assert session_key != "default"
+        uuid.UUID(session_key)
+        assert query_kwargs["parameters"] == [
+            {"name": "@session_id", "value": session_key},
+            {"name": "@source_id", "value": "mem"},
+        ]
+        assert "Received empty session_id" in caplog.text
+
+    async def test_skips_non_dict_message_payload(self, mock_container: MagicMock) -> None:
+        mock_container.query_items.return_value = _to_async_iter([{"message": "bad"}, {"message": None}])
+
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        messages = await provider.get_messages("s1")
+
+        assert messages == []
+
+
+class TestCosmosHistoryProviderListSessions:
+    async def test_list_sessions_returns_unique_sorted_ids(self, mock_container: MagicMock) -> None:
+        mock_container.query_items.return_value = _to_async_iter(["s2", "s1", "s1", "s3"])
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+
+        sessions = await provider.list_sessions()
+
+        assert sessions == ["s1", "s2", "s3"]
+        kwargs = mock_container.query_items.call_args.kwargs
+        assert kwargs["query"] == "SELECT DISTINCT VALUE c.session_id FROM c WHERE c.source_id = @source_id"
+        assert kwargs["parameters"] == [{"name": "@source_id", "value": "mem"}]
+
+
+class TestCosmosHistoryProviderSaveMessages:
+    async def test_saves_messages(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        messages = [Message(role="user", contents=["Hello"]), Message(role="assistant", contents=["Hi"])]
+
+        await provider.save_messages("s1", messages)
+
+        mock_container.execute_item_batch.assert_awaited_once()
+        batch_operations = mock_container.execute_item_batch.await_args.kwargs["batch_operations"]
+        assert len(batch_operations) == 2
+        first_operation, first_args = batch_operations[0]
+        assert first_operation == "upsert"
+        first_document = first_args[0]
+        assert first_document["session_id"] == "s1"
+        assert first_document["message"]["role"] == "user"
+        assert mock_container.execute_item_batch.await_args.kwargs["partition_key"] == "s1"
+
+    async def test_empty_messages_noop(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+
+        await provider.save_messages("s1", [])
+
+        mock_container.execute_item_batch.assert_not_awaited()
+
+    async def test_batches_when_message_count_exceeds_limit(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        messages = [Message(role="user", contents=[f"msg-{index}"]) for index in range(101)]
+
+        await provider.save_messages("s1", messages)
+
+        assert mock_container.execute_item_batch.await_count == 2
+        first_call = mock_container.execute_item_batch.await_args_list[0].kwargs
+        second_call = mock_container.execute_item_batch.await_args_list[1].kwargs
+        assert len(first_call["batch_operations"]) == 100
+        assert len(second_call["batch_operations"]) == 1
+        assert first_call["partition_key"] == "s1"
+        assert second_call["partition_key"] == "s1"
+
+
+class TestCosmosHistoryProviderClear:
+    async def test_clear_deletes_all_session_items(self, mock_container: MagicMock) -> None:
+        mock_container.query_items.return_value = _to_async_iter([{"id": "1"}, {"id": "2"}])
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+
+        await provider.clear("s1")
+
+        mock_container.execute_item_batch.assert_awaited_once()
+        batch_operations = mock_container.execute_item_batch.await_args.kwargs["batch_operations"]
+        assert len(batch_operations) == 2
+        assert batch_operations[0] == ("delete", ("1",))
+        assert batch_operations[1] == ("delete", ("2",))
+        assert mock_container.execute_item_batch.await_args.kwargs["partition_key"] == "s1"
+        query_kwargs = mock_container.query_items.call_args.kwargs
+        assert query_kwargs["query"] == (
+            "SELECT c.id FROM c WHERE c.session_id = @session_id AND c.source_id = @source_id"
+        )
+        assert query_kwargs["parameters"] == [
+            {"name": "@session_id", "value": "s1"},
+            {"name": "@source_id", "value": "mem"},
+        ]
+
+
+class TestCosmosHistoryProviderBeforeAfterRun:
+    async def test_before_run_loads_history(self, mock_container: MagicMock) -> None:
+        msg = Message(role="user", contents=["old msg"])
+        mock_container.query_items.return_value = _to_async_iter([{"message": msg.to_dict()}])
+
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        session = AgentSession(session_id="test")
+        context = SessionContext(input_messages=[Message(role="user", contents=["new msg"])], session_id="s1")
+
+        await provider.before_run(
+            agent=cast(Any, None),
+            session=session,
+            context=context,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+
+        assert "mem" in context.context_messages
+        assert context.context_messages["mem"][0].text == "old msg"
+
+    async def test_after_run_stores_input_and_response(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+        session = AgentSession(session_id="test")
+        context = SessionContext(input_messages=[Message(role="user", contents=["hi"])], session_id="s1")
+        context._response = AgentResponse(messages=[Message(role="assistant", contents=["hello"])])
+
+        await provider.after_run(
+            agent=cast(Any, None),
+            session=session,
+            context=context,
+            state=session.state.setdefault(provider.source_id, {}),
+        )
+
+        mock_container.execute_item_batch.assert_awaited_once()
+        batch_operations = mock_container.execute_item_batch.await_args.kwargs["batch_operations"]
+        assert len(batch_operations) == 2
+        input_doc = batch_operations[0][1][0]
+        response_doc = batch_operations[1][1][0]
+        assert input_doc["message"]["role"] == "user"
+        assert input_doc["message"]["contents"][0]["text"] == "hi"
+        assert response_doc["message"]["role"] == "assistant"
+        assert response_doc["message"]["contents"][0]["text"] == "hello"
+
+
+class TestCosmosHistoryProviderClose:
+    async def test_close_closes_owned_client(
+        self, monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock
+    ) -> None:
+        mock_factory = MagicMock(return_value=mock_cosmos_client)
+        monkeypatch.setattr(history_provider_module, "CosmosClient", mock_factory)
+
+        provider = CosmosHistoryProvider(
+            endpoint="https://account.documents.azure.com:443/",
+            credential="key-123",
+            database_name="db1",
+            container_name="history",
+        )
+
+        await provider.close()
+
+        mock_cosmos_client.close.assert_awaited_once()
+
+    async def test_close_does_not_close_external_client(self, mock_cosmos_client: MagicMock) -> None:
+        provider = CosmosHistoryProvider(
+            source_id="mem",
+            cosmos_client=mock_cosmos_client,
+            database_name="db1",
+            container_name="history",
+        )
+
+        await provider.close()
+
+        mock_cosmos_client.close.assert_not_awaited()
+
+    async def test_async_context_manager_closes_owned_client(
+        self, monkeypatch: pytest.MonkeyPatch, mock_cosmos_client: MagicMock
+    ) -> None:
+        mock_factory = MagicMock(return_value=mock_cosmos_client)
+        monkeypatch.setattr(history_provider_module, "CosmosClient", mock_factory)
+
+        async with CosmosHistoryProvider(
+            endpoint="https://account.documents.azure.com:443/",
+            credential="key-123",
+            database_name="db1",
+            container_name="history",
+        ) as provider:
+            assert provider is not None
+
+        mock_cosmos_client.close.assert_awaited_once()
+
+    async def test_async_context_manager_preserves_original_exception(self, mock_container: MagicMock) -> None:
+        provider = CosmosHistoryProvider(source_id="mem", container_client=mock_container)
+
+        with (
+            patch.object(provider, "close", AsyncMock(side_effect=RuntimeError("close failed"))),
+            pytest.raises(ValueError, match="inner error"),
+        ):
+            async with provider:
+                raise ValueError("inner error")
+
+
+@pytest.mark.flaky
+@pytest.mark.integration
+@skip_if_cosmos_integration_tests_disabled
+async def test_cosmos_history_provider_roundtrip_with_emulator() -> None:
+    endpoint = os.getenv("AZURE_COSMOS_ENDPOINT", "")
+    key = os.getenv("AZURE_COSMOS_KEY", "")
+    database_prefix = os.getenv("AZURE_COSMOS_DATABASE_NAME", "")
+    container_prefix = os.getenv("AZURE_COSMOS_CONTAINER_NAME", "")
+    unique = uuid.uuid4().hex[:8]
+    database_name = f"{database_prefix}-{unique}"
+    container_name = f"{container_prefix}-{unique}"
+    session_id = f"session-{unique}"
+
+    async with CosmosClient(url=endpoint, credential=key) as cosmos_client:
+        await cosmos_client.create_database_if_not_exists(id=database_name)
+        provider = CosmosHistoryProvider(
+            source_id="cosmos_integration",
+            cosmos_client=cosmos_client,
+            database_name=database_name,
+            container_name=container_name,
+        )
+
+        try:
+            await provider.save_messages(
+                session_id,
+                [
+                    Message(role="user", contents=["Hello Cosmos"]),
+                    Message(role="assistant", contents=["Hi from Cosmos"]),
+                ],
+            )
+
+            stored_messages = await provider.get_messages(session_id)
+            assert [message.role for message in stored_messages] == ["user", "assistant"]
+            assert [message.text for message in stored_messages] == ["Hello Cosmos", "Hi from Cosmos"]
+
+            sessions = await provider.list_sessions()
+            assert session_id in sessions
+
+            await provider.clear(session_id)
+            assert await provider.get_messages(session_id) == []
+        finally:
+            with suppress(CosmosResourceNotFoundError):
+                await cosmos_client.delete_database(database_name)