diff --git a/src/agents/extensions/memory/dapr_session.py b/src/agents/extensions/memory/dapr_session.py index 20b300ce3c..f9462fd465 100644 --- a/src/agents/extensions/memory/dapr_session.py +++ b/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. diff --git a/tests/extensions/memory/test_dapr_session.py b/tests/extensions/memory/test_dapr_session.py index 702cb7388c..5458d39128 100644 --- a/tests/extensions/memory/test_dapr_session.py +++ b/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()