From 93a346392d3cffcf9e1e2581770910492eeb57c9 Mon Sep 17 00:00:00 2001
From: Kazuhiro Sera <seratch@openai.com>
Date: Thu, 28 May 2026 15:08:54 +0900
Subject: [PATCH] fix review comments

---
 src/agents/run_context.py           | 14 ++++++++++++-
 tests/test_run_context_approvals.py | 32 +++++++++++++++++++++++++++++
 tests/test_run_state.py             | 15 ++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/src/agents/run_context.py b/src/agents/run_context.py
index b7f03038..1916748c 100644
--- a/src/agents/run_context.py
+++ b/src/agents/run_context.py
@@ -3,7 +3,7 @@ from __future__ import annotations
 from collections.abc import Mapping
 from dataclasses import dataclass, field
 from threading import RLock
-from typing import TYPE_CHECKING, Any, Generic
+from typing import TYPE_CHECKING, Any, Generic, SupportsIndex
 
 from typing_extensions import TypeVar
 
@@ -48,6 +48,9 @@ class _ApprovalRecords(dict[str, _ApprovalRecord]):
         super().__init__(*args, **kwargs)
         self.lock = RLock() if lock is None else lock
 
+    def __reduce_ex__(self, protocol: SupportsIndex) -> tuple[Any, ...]:
+        return (self.__class__, (dict(self),))
+
 
 @dataclass(eq=False)
 class RunContextWrapper(Generic[TContext]):
@@ -78,6 +81,15 @@ class RunContextWrapper(Generic[TContext]):
             self._approvals = _ApprovalRecords(self._approvals)
             self._approvals_lock = self._approvals.lock
 
+    def __getstate__(self) -> dict[str, Any]:
+        state = self.__dict__.copy()
+        state.pop("_approvals_lock", None)
+        return state
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        self.__dict__.update(state)
+        self.__post_init__()
+
     @staticmethod
     def _to_str_or_none(value: Any) -> str | None:
         if isinstance(value, str):
diff --git a/tests/test_run_context_approvals.py b/tests/test_run_context_approvals.py
index 674fe2ba..716953cc 100644
--- a/tests/test_run_context_approvals.py
+++ b/tests/test_run_context_approvals.py
@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+import copy
+import pickle
+from typing import Any
+
 from agents import Agent, RunContextWrapper
 from agents.run_context import AgentHookContext
 from agents.tool_context import ToolContext
@@ -262,3 +266,31 @@ def test_contexts_constructed_with_shared_approvals_reuse_approval_lock() -> Non
     assert tool_context._approvals_lock is context_wrapper._approvals_lock
     assert hook_context._approvals is context_wrapper._approvals
     assert hook_context._approvals_lock is context_wrapper._approvals_lock
+
+
+def test_context_approval_lock_is_recreated_after_deepcopy() -> None:
+    agent = Agent(name="test-agent")
+    context_wrapper = RunContextWrapper(context=None)
+    approval_item = make_tool_approval_item(agent, call_id="call-1", name="lookup")
+
+    context_wrapper.approve_tool(approval_item)
+    copied_context = copy.deepcopy(context_wrapper)
+
+    assert copied_context.is_tool_approved("lookup", "call-1") is True
+    assert copied_context._approvals is not context_wrapper._approvals
+    copied_approvals: Any = copied_context._approvals
+    assert copied_context._approvals_lock is copied_approvals.lock
+
+
+def test_context_approval_lock_is_recreated_after_pickle_roundtrip() -> None:
+    agent = Agent(name="test-agent")
+    context_wrapper = RunContextWrapper(context=None)
+    approval_item = make_tool_approval_item(agent, call_id="call-1", name="lookup")
+
+    context_wrapper.approve_tool(approval_item)
+    restored_context = pickle.loads(pickle.dumps(context_wrapper))
+
+    assert restored_context.is_tool_approved("lookup", "call-1") is True
+    assert restored_context._approvals is not context_wrapper._approvals
+    restored_approvals: Any = restored_context._approvals
+    assert restored_context._approvals_lock is restored_approvals.lock
diff --git a/tests/test_run_state.py b/tests/test_run_state.py
index 394efe57..63e55d3a 100644
--- a/tests/test_run_state.py
+++ b/tests/test_run_state.py
@@ -6,6 +6,7 @@ import gc
 import io
 import json
 import logging
+import pickle
 import threading
 import time
 from collections.abc import AsyncIterator, Callable, Mapping
@@ -1620,6 +1621,20 @@ class TestRunState:
         assert "tool-new" not in serialized
         assert context.is_tool_approved("tool-new", "cid-new") is True
 
+    def test_approval_lock_is_recreated_after_run_state_pickle_roundtrip(self):
+        """Test that pickling RunState does not serialize the approval lock."""
+        context: RunContextWrapper[dict[str, str]] = RunContextWrapper(context={})
+        agent = Agent(name="ApprovalAgent")
+        state = make_state(agent, context=context, original_input="test")
+        state.approve(make_tool_approval_item(agent, call_id="cid1", name="tool1"))
+
+        restored_state = pickle.loads(pickle.dumps(state))
+
+        assert restored_state._context is not None
+        assert restored_state._context.is_tool_approved("tool1", "cid1") is True
+        restored_approvals: Any = restored_state._context._approvals
+        assert restored_state._context._approvals_lock is restored_approvals.lock
+
     async def test_serializes_and_restores_rejection_messages(self):
         """Test that rejection messages are preserved through serialization."""
         context: RunContextWrapper[dict[str, str]] = RunContextWrapper(context={})