fix(memory): preserve DaprSession created_at across writes

src/agents/extensions/memory/dapr_session.py

@@ -43,7 +43,11 @@
     )
 
 from ...items import TResponseInputItem
-from ...logger import log_model_and_tool_action_error, logger
+from ...logger import (
+    log_model_and_tool_action_error,
+    log_model_and_tool_action_warning,
+    logger,
+)
 from ...memory.session import SessionABC
 from ...memory.session_settings import (
     SessionSettings,
@@ -184,6 +188,32 @@ def _get_metadata(self) -> dict[str, str]:
             metadata["ttlInSeconds"] = str(self._ttl)
         return metadata
 
+    async def _read_created_at(self) -> tuple[str | None, str | None]:
+        """Return the stored creation timestamp and the metadata etag backing it.
+
+        The timestamp is None when it is missing or unreadable. The etag is returned so the
+        caller can write the metadata conditionally against the same revision it read. The
+        Dapr SDK reports a missing etag as an empty string, so it is normalized to None and
+        callers can test for a real etag rather than for a particular empty representation.
+        """
+        response = await self._dapr_client.get_state(
+            store_name=self._state_store_name,
+            key=self._metadata_key,
+            state_metadata=self._get_read_metadata(),
+        )
+        etag = response.etag or None
+        data = response.data
+        if not data:
+            return None, etag
+        try:
+            stored = json.loads(data.decode("utf-8"))
+        except (json.JSONDecodeError, UnicodeDecodeError):
+            return None, etag
+        if not isinstance(stored, dict):
+            return None, etag
+        created_at = stored.get("created_at")
+        return (created_at if isinstance(created_at, str) and created_at else None), etag
+
     async def _serialize_item(self, item: TResponseInputItem) -> str:
         """Serialize an item to JSON string. Can be overridden by subclasses."""
         return json.dumps(item, separators=(",", ":"))
@@ -362,19 +392,55 @@ async def add_items(self, items: list[TResponseInputItem]) -> None:
                         continue
                     raise
 
-            # Update metadata
-            metadata = {
-                "session_id": self.session_id,
-                "created_at": str(int(time.time())),
-                "updated_at": str(int(time.time())),
-            }
-            await self._dapr_client.save_state(
-                store_name=self._state_store_name,
-                key=self._metadata_key,
-                value=json.dumps(metadata),
-                state_metadata=self._get_metadata(),
-                options=self._get_state_options(),
-            )
+            # Update metadata, preserving created_at across subsequent writes. A plain write
+            # would let a later append overwrite created_at with its own now, so the save is
+            # guarded by the etag that backed the value that was read.
+            #
+            # The guard only applies once a real etag was read. Dapr documents a write without
+            # an etag as last-write-wins even when first-write concurrency is requested, so
+            # asking for it on the create is not a race guarantee and is left off rather than
+            # implying one. Concurrent etag-less creation is therefore last-write-wins.
+            #
+            # The messages key is already committed once we reach here, so raising would tell
+            # the caller the append failed while their items are in the session, and the
+            # natural response of retrying add_items() would store the batch twice. Metadata
+            # is derived bookkeeping, so give up on it with a warning instead.
+            attempt = 0
+            while True:
+                attempt += 1
+                try:
+                    stored_created_at, metadata_etag = await self._read_created_at()
+                    now = str(int(time.time()))
+                    metadata = {
+                        "session_id": self.session_id,
+                        "created_at": stored_created_at or now,
+                        "updated_at": now,
+                    }
+                    await self._dapr_client.save_state(
+                        store_name=self._state_store_name,
+                        key=self._metadata_key,
+                        value=json.dumps(metadata),
+                        etag=metadata_etag,
+                        state_metadata=self._get_metadata(),
+                        options=self._get_state_options(
+                            concurrency=(
+                                Concurrency.first_write if metadata_etag is not None else None
+                            )
+                        ),
+                    )
+                    break
+                except Exception as error:
+                    should_retry = await self._handle_concurrency_conflict(error, attempt)
+                    if should_retry:
+                        continue
+                    log_model_and_tool_action_warning(
+                        logger,
+                        "DaprSession stored the new items but could not update the session "
+                        "metadata",
+                        error,
+                        diagnostic_extra=lambda: {"session_id": self.session_id},
+                    )
+                    break
 
     async def pop_item(self) -> TResponseInputItem | None:
         """Remove and return the most recent item from the session.

tests/extensions/memory/test_dapr_session.py

@@ -40,7 +40,8 @@ async def get_state(
         """Get state from in-memory store."""
         response = Mock()
         response.data = self._state.get(key, b"")
-        response.etag = self._etags.get(key)
+        # The Dapr SDK reports a missing etag as an empty string rather than None.
+        response.etag = self._etags.get(key, "")
         return response
 
     async def save_state(
@@ -1327,3 +1328,263 @@ def entry_signalling_resolve(*args: Any, **kwargs: Any) -> Any:
                 task.cancel()
                 with suppress(asyncio.CancelledError, RuntimeError):
                     await task
+
+
+async def test_add_items_preserves_created_at_metadata(
+    fake_dapr_client: FakeDaprClient, monkeypatch: pytest.MonkeyPatch
+):
+    """`created_at` must be set once and not overwritten by subsequent add_items calls."""
+    session = await _create_test_session(fake_dapr_client)
+
+    try:
+        monkeypatch.setattr("agents.extensions.memory.dapr_session.time.time", lambda: 1000.0)
+        await session.add_items([{"role": "user", "content": "first"}])
+        first = json.loads(fake_dapr_client._state[session._metadata_key].decode("utf-8"))
+        assert first["created_at"] == "1000"
+        assert first["updated_at"] == "1000"
+
+        monkeypatch.setattr("agents.extensions.memory.dapr_session.time.time", lambda: 2000.0)
+        await session.add_items([{"role": "user", "content": "second"}])
+        second = json.loads(fake_dapr_client._state[session._metadata_key].decode("utf-8"))
+        assert second["created_at"] == "1000"
+        assert second["updated_at"] == "2000"
+    finally:
+        await session.close()
+
+
+async def test_metadata_creation_does_not_request_first_write(
+    fake_dapr_client: FakeDaprClient, monkeypatch: pytest.MonkeyPatch
+):
+    """Creating the metadata key must not ask for first-write concurrency.
+
+    Dapr treats a write with no etag as last-write-wins even when first-write is requested,
+    so asking for it on the create would imply a guarantee the store does not provide. The
+    SDK reports a missing etag as an empty string, so this also pins that the empty value is
+    not mistaken for a real one.
+    """
+    session = await _create_test_session(fake_dapr_client, "metadata_create_concurrency")
+
+    try:
+        real_save = fake_dapr_client.save_state
+        seen: list[tuple[str | None, Any]] = []
+
+        async def record_metadata_saves(
+            store_name: str,
+            key: str,
+            value: str | bytes,
+            **kwargs: Any,
+        ) -> None:
+            if key == session._metadata_key:
+                seen.append(
+                    (kwargs.get("etag"), getattr(kwargs.get("options"), "concurrency", None))
+                )
+            await real_save(store_name, key, value, **kwargs)
+
+        monkeypatch.setattr(fake_dapr_client, "save_state", record_metadata_saves)
+
+        await session.add_items([{"role": "user", "content": "first"}])
+        assert len(seen) == 1
+        create_etag, create_concurrency = seen[0]
+        # No real etag existed, so no etag is sent and concurrency is left unspecified
+        # rather than first-write, which Dapr would ignore here anyway.
+        assert create_etag is None
+        assert getattr(create_concurrency, "name", None) == "unspecified"
+
+        await session.add_items([{"role": "user", "content": "second"}])
+        assert len(seen) == 2
+        update_etag, update_concurrency = seen[1]
+        # Metadata now exists, so the update is guarded by the etag that backed it.
+        assert update_etag is not None
+        assert getattr(update_concurrency, "name", None) == "first_write"
+    finally:
+        await session.close()
+
+
+async def test_stale_metadata_etag_retries_and_keeps_created_at(
+    fake_dapr_client: FakeDaprClient, monkeypatch: pytest.MonkeyPatch
+):
+    """A metadata save guarded by a stale etag must retry and keep the stored `created_at`.
+
+    Once the metadata key exists, every later append reads a real etag and saves against it.
+    A writer holding an etag that another append has already superseded is rejected, so it
+    re-reads and adopts the stored `created_at` rather than replacing it with its own `now`.
+
+    Dapr treats a write with no etag as last-write-wins even when first-write concurrency is
+    requested, so the create is deliberately not covered here. This is the guarantee the
+    store actually provides.
+    """
+    session_id = "shared_session_created_at_stale_etag"
+    first = await _create_test_session(fake_dapr_client, session_id=session_id)
+    second = await _create_test_session(fake_dapr_client, session_id=session_id)
+
+    try:
+        monkeypatch.setattr("agents.extensions.memory.dapr_session.time.time", lambda: 1000.0)
+        await first.add_items([{"role": "user", "content": "first"}])
+        created = json.loads(fake_dapr_client._state[first._metadata_key].decode("utf-8"))
+        assert created["created_at"] == "1000"
+        stale_etag = fake_dapr_client._etags[first._metadata_key]
+
+        # A second append supersedes that etag, so the value captured above is now stale but
+        # still non-null, which is the situation the guard is actually for.
+        monkeypatch.setattr("agents.extensions.memory.dapr_session.time.time", lambda: 1500.0)
+        await first.add_items([{"role": "user", "content": "second"}])
+        assert fake_dapr_client._etags[first._metadata_key] != stale_etag
+
+        real_read = second._read_created_at
+        calls = 0
+
+        async def read_stale_once() -> tuple[str | None, str | None]:
+            nonlocal calls
+            calls += 1
+            if calls == 1:
+                # Pretend this writer read the metadata before the second append landed.
+                return "1000", stale_etag
+            return await real_read()
+
+        monkeypatch.setattr(second, "_read_created_at", read_stale_once)
+        monkeypatch.setattr("agents.extensions.memory.dapr_session.time.time", lambda: 2000.0)
+        await second.add_items([{"role": "user", "content": "third"}])
+
+        final = json.loads(fake_dapr_client._state[second._metadata_key].decode("utf-8"))
+        # The stale save is rejected, so the retry reads the stored value and keeps it.
+        assert final["created_at"] == "1000"
+        assert final["updated_at"] == "2000"
+        # Two reads means the conditional save actually rejected the stale one and retried.
+        assert calls == 2
+    finally:
+        await first.close()
+        await second.close()
+
+
+async def test_add_items_survives_metadata_write_giving_up(
+    fake_dapr_client: FakeDaprClient, monkeypatch: pytest.MonkeyPatch, caplog: Any
+):
+    """A metadata write that exhausts its retries must not fail an append that already landed.
+
+    The messages key is saved before the metadata key, so raising here would report failure for
+    items that are already in the session, and the natural retry of `add_items` would store the
+    same batch a second time.
+    """
+    import logging
+
+    session = await _create_test_session(fake_dapr_client, "metadata_gives_up")
+
+    try:
+        real_save = fake_dapr_client.save_state
+
+        async def fail_only_metadata_saves(
+            store_name: str,
+            key: str,
+            value: str | bytes,
+            **kwargs: Any,
+        ) -> None:
+            if key == session._metadata_key:
+                raise RuntimeError("etag mismatch")
+            await real_save(store_name, key, value, **kwargs)
+
+        monkeypatch.setattr(fake_dapr_client, "save_state", fail_only_metadata_saves)
+        monkeypatch.setattr(session, "_calculate_retry_delay", lambda attempt: 0.0)
+
+        with caplog.at_level(logging.WARNING):
+            await session.add_items([{"role": "user", "content": "kept"}])
+
+        items = await session.get_items()
+        assert [item.get("content") for item in items] == ["kept"]
+
+        warnings = [
+            record for record in caplog.records if "could not update" in record.getMessage()
+        ]
+        assert warnings
+        # Data logging is off by default, so neither the caller-supplied session id nor the
+        # provider error text may reach the record.
+        for record in warnings:
+            assert "metadata_gives_up" not in record.getMessage()
+            assert "etag mismatch" not in record.getMessage()
+    finally:
+        await session.close()
+
+
+@pytest.mark.parametrize(
+    ("dont_log_model_data", "dont_log_tool_data", "redacted"),
+    [
+        (True, False, True),
+        (False, True, True),
+        (False, False, False),
+    ],
+    ids=["model-redacted", "tool-redacted", "fully-diagnostic"],
+)
+async def test_metadata_write_warning_respects_data_policies(
+    fake_dapr_client: FakeDaprClient,
+    monkeypatch: pytest.MonkeyPatch,
+    caplog: Any,
+    dont_log_model_data: bool,
+    dont_log_tool_data: bool,
+    redacted: bool,
+):
+    """The give-up warning must obey both data policies, and never lose the committed items.
+
+    `log_model_and_tool_action_warning` redacts when either policy is on, so only the mode with
+    both off may carry the session id or the provider error. This inspects the whole LogRecord
+    rather than just the rendered message, because the session id travels in `extra` and the
+    exception travels in `exc_info`, neither of which shows up in `getMessage()`.
+    """
+    import logging
+
+    import agents._debug as _debug
+
+    session_id = f"metadata_policy_{dont_log_model_data}_{dont_log_tool_data}"
+    session = await _create_test_session(fake_dapr_client, session_id)
+
+    try:
+        monkeypatch.setattr(_debug, "DONT_LOG_MODEL_DATA", dont_log_model_data)
+        monkeypatch.setattr(_debug, "DONT_LOG_TOOL_DATA", dont_log_tool_data)
+
+        real_save = fake_dapr_client.save_state
+        error_text = "etag mismatch"
+
+        async def fail_only_metadata_saves(
+            store_name: str,
+            key: str,
+            value: str | bytes,
+            **kwargs: Any,
+        ) -> None:
+            if key == session._metadata_key:
+                raise RuntimeError(error_text)
+            await real_save(store_name, key, value, **kwargs)
+
+        monkeypatch.setattr(fake_dapr_client, "save_state", fail_only_metadata_saves)
+        monkeypatch.setattr(session, "_calculate_retry_delay", lambda attempt: 0.0)
+
+        with caplog.at_level(logging.WARNING):
+            await session.add_items([{"role": "user", "content": "kept"}])
+
+        # The append landed regardless of how the failure was logged.
+        items = await session.get_items()
+        assert [item.get("content") for item in items] == ["kept"]
+
+        records = [record for record in caplog.records if "could not update" in record.getMessage()]
+        assert len(records) == 1
+        record = records[0]
+        rendered = logging.Formatter().format(record)
+
+        if redacted:
+            # Redacted form is the bare message with no exception and no diagnostic context.
+            assert record.msg == "%s"
+            assert record.exc_info is None
+            assert record.exc_text is None
+            assert not hasattr(record, "openai_agents_diagnostic_context")
+            assert session_id not in rendered
+            assert error_text not in rendered
+            assert all(
+                session_id not in str(value) and error_text not in str(value)
+                for value in record.__dict__.values()
+            )
+        else:
+            # Diagnostic form carries the exception and the session id, by design.
+            assert record.msg == "%s: %s"
+            assert record.exc_info is not None
+            assert isinstance(record.exc_info[1], RuntimeError)
+            assert error_text in rendered
+            assert record.openai_agents_diagnostic_context == {"session_id": session_id}
+    finally:
+        await session.close()