fix review comments

Author: seratch

src/agents/memory/openai_responses_compaction_session.py

@@ -186,17 +186,58 @@ def _resolve_compaction_mode(
 
         return resolved_mode
 
+    def _resolve_store_tracking(
+        self,
+        *,
+        response_id: str | None,
+        previous_response_id: str | None,
+        store: bool | None,
+        store_was_provided: bool,
+    ) -> bool | None:
+        """Resolve the effective store setting for the current response id.
+
+        Reuse `_last_store` only while compaction still refers to the same response. A new
+        response id with no explicit `store` falls back to the Responses API default behavior.
+        """
+        if store_was_provided:
+            self._last_store = store
+            return store
+
+        if response_id is not None and response_id != previous_response_id:
+            self._last_store = None
+            return None
+
+        return self._last_store
+
+    def _get_effective_store_for_response_id(
+        self,
+        *,
+        response_id: str | None,
+        store: bool | None,
+    ) -> bool | None:
+        """Return the effective store setting without mutating response tracking."""
+        if store is not None:
+            return store
+        if response_id is not None and response_id != self._response_id:
+            return None
+        return self._last_store
+
     async def run_compaction(self, args: OpenAIResponsesCompactionArgs | None = None) -> None:
         """Run compaction using responses.compact API."""
         previous_response_id = self._response_id
         if args and args.get("response_id"):
             self._response_id = args["response_id"]
         requested_mode = args.get("compaction_mode") if args else None
-        if args and "store" in args:
-            store: bool | None = args["store"]
-            self._last_store = store
-        else:
-            store = self._last_store
+        store_was_provided = bool(args and "store" in args)
+        requested_store: bool | None = (
+            args["store"] if args is not None and "store" in args else None
+        )
+        store = self._resolve_store_tracking(
+            response_id=self._response_id,
+            previous_response_id=previous_response_id,
+            store=requested_store,
+            store_was_provided=store_was_provided,
+        )
         turn_has_local_adds_without_new_response_id = (
             self._has_unacknowledged_local_session_adds
             and (args is None or args.get("response_id") in {None, previous_response_id})
@@ -323,7 +364,7 @@ async def _defer_compaction(self, response_id: str, store: bool | None = None) -
         resolved_mode = _resolve_compaction_mode(
             self.compaction_mode,
             response_id=response_id,
-            store=store if store is not None else self._last_store,
+            store=self._get_effective_store_for_response_id(response_id=response_id, store=store),
         )
         should_compact = self.should_trigger_compaction(
             {

src/agents/run.py

@@ -47,6 +47,7 @@
 from .run_internal.agent_runner_helpers import (
     append_model_response_if_new,
     apply_resumed_conversation_settings,
+    attach_run_state_metadata,
     build_interruption_result,
     build_resumed_stream_debug_extra,
     ensure_context_wrapper,
@@ -828,8 +829,6 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                                 result._replay_from_model_input_items = list(
                                     generated_items
                                 ) != list(session_items)
-                                if run_state is not None:
-                                    result._trace_state = run_state._trace_state
                                 if session_persistence_enabled:
                                     input_items_for_save_1: list[TResponseInputItem] = (
                                         session_input_items_for_persistence
@@ -844,6 +843,7 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                                         response_id=turn_result.model_response.response_id,
                                         store=store_setting,
                                     )
+                                attach_run_state_metadata(result, run_state=run_state)
                                 result._original_input = copy_input_items(original_input)
                                 return finalize_conversation_tracking(
                                     _with_reasoning_item_id_policy(result),
@@ -965,8 +965,6 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                         result._replay_from_model_input_items = list(generated_items) != list(
                             session_items
                         )
-                        if run_state is not None:
-                            result._trace_state = run_state._trace_state
                         if session_persistence_enabled and include_in_history:
                             handler_input_items_for_save: list[TResponseInputItem] = (
                                 session_input_items_for_persistence
@@ -981,6 +979,7 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                                 response_id=None,
                                 store=store_setting,
                             )
+                        attach_run_state_metadata(result, run_state=run_state)
                         result._original_input = copy_input_items(original_input)
                         return finalize_conversation_tracking(
                             _with_reasoning_item_id_policy(result),
@@ -1236,10 +1235,6 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                             result._replay_from_model_input_items = list(generated_items) != list(
                                 session_items
                             )
-                            if run_state is not None:
-                                result._current_turn_persisted_item_count = (
-                                    run_state._current_turn_persisted_item_count
-                                )
                             await save_turn_items_if_needed(
                                 session=session,
                                 run_state=run_state,
@@ -1249,6 +1244,7 @@ def _with_reasoning_item_id_policy(result: RunResult) -> RunResult:
                                 response_id=turn_result.model_response.response_id,
                                 store=store_setting,
                             )
+                            attach_run_state_metadata(result, run_state=run_state)
                             result._original_input = copy_input_items(original_input)
                             return finalize_conversation_tracking(
                                 _with_reasoning_item_id_policy(result),

src/agents/run_internal/agent_runner_helpers.py

@@ -33,6 +33,7 @@
 from .tool_use_tracker import AgentToolUseTracker, serialize_tool_use_tracker
 
 __all__ = [
+    "attach_run_state_metadata",
     "apply_resumed_conversation_settings",
     "append_model_response_if_new",
     "build_generated_items_details",
@@ -276,6 +277,17 @@ def finalize_conversation_tracking(
     return result
 
 
+def attach_run_state_metadata(result: RunResult, *, run_state: RunState | None) -> RunResult:
+    """Copy resumable state metadata from the current RunState onto a RunResult."""
+    if run_state is None:
+        return result
+
+    result._state = run_state
+    result._current_turn_persisted_item_count = run_state._current_turn_persisted_item_count
+    result._trace_state = run_state._trace_state
+    return result
+
+
 def build_interruption_result(
     *,
     result_input: str | list[TResponseInputItem],
@@ -315,10 +327,7 @@ def build_interruption_result(
     result._current_turn = current_turn
     result._model_input_items = list(generated_items)
     result._replay_from_model_input_items = list(generated_items) != list(session_items)
-    if run_state is not None:
-        result._state = run_state
-        result._current_turn_persisted_item_count = run_state._current_turn_persisted_item_count
-        result._trace_state = run_state._trace_state
+    attach_run_state_metadata(result, run_state=run_state)
     result._original_input = copy_input_items(original_input)
     return result
 

tests/memory/test_openai_responses_compaction_session.py

@@ -551,8 +551,8 @@ async def test_run_compaction_auto_uses_default_store_when_unset(self) -> None:
         first_kwargs = mock_client.responses.compact.call_args_list[0].kwargs
         second_kwargs = mock_client.responses.compact.call_args_list[1].kwargs
         assert "previous_response_id" not in first_kwargs
-        assert "previous_response_id" not in second_kwargs
-        assert second_kwargs.get("input") == []
+        assert second_kwargs.get("previous_response_id") == "resp-stored"
+        assert "input" not in second_kwargs
 
     @pytest.mark.asyncio
     async def test_run_compaction_auto_uses_input_when_last_response_unstored(self) -> None:

tests/test_agent_tracing.py

@@ -341,6 +341,23 @@ def send_email(recipient: str) -> str:
     assert len(after_custom_spans) == len(before_custom_spans)
 
 
+@pytest.mark.asyncio
+async def test_completed_result_to_state_preserves_sensitive_trace_flag() -> None:
+    model = FakeModel()
+    model.add_multiple_turn_outputs([[get_text_message("done")]])
+    agent = Agent(name="trace_agent", model=model)
+
+    result = await Runner.run(
+        agent,
+        input="first_test",
+        run_config=RunConfig(trace_include_sensitive_data=False),
+    )
+
+    state = result.to_state()
+
+    assert state._trace_include_sensitive_data is False
+
+
 @pytest.mark.asyncio
 async def test_wrapped_trace_is_single_trace():
     model = FakeModel()