python-codex 0.1.8__py3-none-any.whl → 0.1.10__py3-none-any.whl

pycodex/context.py

@@ -89,7 +89,7 @@ class ContextConfig:
         profile: 'typing.Union[str, None]' = None,
     ) -> 'ContextConfig':
         path = Path(config_path)
-        data = tomllib.loads(path.read_text())
+        data = tomllib.loads(path.read_text(encoding="utf-8"))
         selected = dict(data)
         if profile is not None:
             overrides = data.get("profiles", {}).get(profile)
@@ -162,7 +162,9 @@ class ContextManager:
         self._include_permissions_instructions = include_permissions_instructions
         self._include_skills_instructions = include_skills_instructions
         self._network_access = network_access
-        self._default_base_instructions = DEFAULT_BASE_INSTRUCTIONS_PATH.read_text()
+        self._default_base_instructions = DEFAULT_BASE_INSTRUCTIONS_PATH.read_text(
+            encoding="utf-8"
+        )
         self._workspace_metadata_turn_id: 'typing.Union[str, None]' = None
         self._workspace_metadata_cache: 'typing.Union[JSONDict, None]' = None
 
@@ -237,7 +239,10 @@ class ContextManager:
         if self._config.base_instructions is not None:
             return self._config.base_instructions
         if self._config.model_instructions_file is not None:
-            return self._config.model_instructions_file.read_text().strip()
+            return self._config.model_instructions_file.read_text(
+                encoding="utf-8",
+                errors="replace",
+            ).strip()
         resolved = self._resolve_model_instructions()
         if resolved is not None:
             return resolved
@@ -327,11 +332,11 @@ class ContextManager:
             return None
 
         sandbox_text = (
-            sandbox_prompt_path.read_text().strip().replace(
+            sandbox_prompt_path.read_text(encoding="utf-8").strip().replace(
                 "{network_access}", self._network_access
             )
         )
-        approval_text = approval_prompt_path.read_text().strip()
+        approval_text = approval_prompt_path.read_text(encoding="utf-8").strip()
         return "\n".join(
             [
                 PERMISSIONS_OPEN_TAG,
@@ -429,7 +434,7 @@ class ContextManager:
         docs: 'typing.List[str]' = []
         remaining = self._config.project_doc_max_bytes
         for path in self._discover_project_doc_paths():
-            text = path.read_text()
+            text = path.read_text(encoding="utf-8", errors="replace")
             if not text.strip():
                 continue
             if remaining is None:
@@ -437,7 +442,7 @@ class ContextManager:
                 continue
             if remaining <= 0:
                 break
-            encoded = text.encode()
+            encoded = text.encode("utf-8")
             docs.append(encoded[:remaining].decode(errors="ignore"))
             remaining -= min(len(encoded), remaining)
         if not docs:
@@ -507,15 +512,15 @@ def _normalize_int(value) -> 'typing.Union[int, None]':
 
 def _default_collaboration_instructions(mode: 'CollaborationMode') -> 'str':
     if mode == "plan":
-        return PLAN_COLLABORATION_INSTRUCTIONS_PATH.read_text()
-    return DEFAULT_COLLABORATION_INSTRUCTIONS_PATH.read_text()
+        return PLAN_COLLABORATION_INSTRUCTIONS_PATH.read_text(encoding="utf-8")
+    return DEFAULT_COLLABORATION_INSTRUCTIONS_PATH.read_text(encoding="utf-8")
 
 
 def _read_first_instruction_file(base: 'Path') -> 'typing.Union[str, None]':
     for candidate_name in (LOCAL_PROJECT_DOC_FILENAME, DEFAULT_PROJECT_DOC_FILENAME):
         candidate = base / candidate_name
         try:
-            contents = candidate.read_text()
+            contents = candidate.read_text(encoding="utf-8", errors="replace")
         except OSError:
             continue
         trimmed = contents.strip()
@@ -526,7 +531,7 @@ def _read_first_instruction_file(base: 'Path') -> 'typing.Union[str, None]':
 
 @lru_cache(maxsize=1)
 def _load_models_by_slug() -> 'typing.Dict[str, JSONDict]':
-    payload = json.loads(DEFAULT_MODELS_PATH.read_text())
+    payload = json.loads(DEFAULT_MODELS_PATH.read_text(encoding="utf-8"))
     models = payload.get("models", [])
     by_slug: 'typing.Dict[str, JSONDict]' = {}
     for model in models:
@@ -571,7 +576,7 @@ def _discover_skill_files(
 
 
 def _parse_skill_descriptor(path: 'Path', scope_rank: 'int') -> 'typing.Union[SkillDescriptor, None]':
-    text = path.read_text()
+    text = path.read_text(encoding="utf-8", errors="replace")
     if not text.startswith("---\n"):
         return None
     end_marker = "\n---\n"

pycodex/model.py

@@ -71,7 +71,7 @@ class ResponsesProviderConfig:
         config_path: 'typing.Union[str, Path]' = DEFAULT_CODEX_CONFIG_PATH,
         profile: 'typing.Union[str, None]' = None,
     ) -> 'ResponsesProviderConfig':
-        data = tomllib.loads(Path(config_path).read_text())
+        data = tomllib.loads(Path(config_path).read_text(encoding="utf-8"))
         selected = dict(data)
         if profile is not None:
             overrides = data.get("profiles", {}).get(profile)

pycodex/portable.py

@@ -123,7 +123,8 @@ def bootstrap_called_home(
             },
             ensure_ascii=False,
             indent=2,
-        )
+        ),
+        encoding="utf-8",
     )
     return home_dir / DEFAULT_ENTRY_CONFIG
 
@@ -199,7 +200,7 @@ def _collect_config_referenced_files(root: 'Path') -> 'typing.Set[str]':
     config_path = root / DEFAULT_ENTRY_CONFIG
     if not config_path.is_file():
         return set()
-    data = tomllib.loads(config_path.read_text())
+    data = tomllib.loads(config_path.read_text(encoding="utf-8"))
     referenced: 'typing.Set[str]' = set()
     candidates = [data]
     profiles = data.get("profiles")
@@ -352,7 +353,7 @@ def _load_cached_metadata(metadata_path: 'Path') -> 'typing.Dict[str, object]':
     if not metadata_path.is_file():
         return {}
     try:
-        payload = json.loads(metadata_path.read_text())
+        payload = json.loads(metadata_path.read_text(encoding="utf-8"))
     except (ValueError, OSError):
         return {}
     return payload if isinstance(payload, dict) else {}

pycodex/tools/base_tool.py

@@ -30,7 +30,8 @@ EXEC_TOOLS_SNAPSHOT_PATH = (
 @lru_cache(maxsize=1)
 def _load_exec_tool_payloads() -> 'typing.Dict[str, JSONDict]':
     payloads: 'typing.Dict[str, JSONDict]' = {}
-    for payload in json.loads(EXEC_TOOLS_SNAPSHOT_PATH.read_text()):
+    raw_payloads = EXEC_TOOLS_SNAPSHOT_PATH.read_text(encoding="utf-8")
+    for payload in json.loads(raw_payloads):
         if not isinstance(payload, dict):
             continue
         name = payload.get("name")

pycodex/utils/dotenv.py

@@ -18,7 +18,9 @@ def load_codex_dotenv(config_path: 'typing.Union[str, Path]') -> 'None':
         _LOADED_CODEX_DOTENV_HOMES.add(codex_home)
         return
 
-    for key, value in parse_dotenv(dotenv_path.read_text()).items():
+    for key, value in parse_dotenv(
+        dotenv_path.read_text(encoding="utf-8", errors="replace")
+    ).items():
         if key.upper().startswith(ILLEGAL_ENV_VAR_PREFIX):
             continue
         os.environ[key] = value

pycodex/utils/get_env.py

@@ -98,7 +98,11 @@ def get_os_info() -> 'typing.Tuple[str, str]':
     os_release = Path("/etc/os-release")
     if os_release.is_file():
         values: 'typing.Dict[str, str]' = {}
-        for line in os_release.read_text().splitlines():
+        os_release_text = os_release.read_text(
+            encoding="utf-8",
+            errors="replace",
+        )
+        for line in os_release_text.splitlines():
             if "=" not in line:
                 continue
             key, value = line.split("=", 1)

pycodex/utils/session_persist.py

@@ -282,7 +282,9 @@ def _latest_thread_names_by_id(codex_home: 'Path') -> 'typing.Dict[str, str]':
         return {}
 
     names_by_id: 'typing.Dict[str, str]' = {}
-    for raw_line in reversed(index_path.read_text().splitlines()):
+    for raw_line in reversed(
+        index_path.read_text(encoding="utf-8", errors="replace").splitlines()
+    ):
         line = raw_line.strip()
         if not line:
             continue
@@ -321,7 +323,7 @@ def _extract_first_user_message_preview(rollout_path: 'Path') -> 'typing.Union[s
 
 
 def _iter_rollout_entries(rollout_path: 'Path') -> 'typing.Iterable[typing.Dict[str, object]]':
-    text = rollout_path.read_text()
+    text = rollout_path.read_text(encoding="utf-8", errors="replace")
     decoder = json.JSONDecoder()
     index = 0
     parsed_entries = 0

{python_codex-0.1.8.dist-info → python_codex-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: python-codex
-Version: 0.1.8
+Version: 0.1.10
 Summary: A minimal Python extraction of Codex's main agent loop
 License-File: LICENSE
 Requires-Python: >=3.6.2

{python_codex-0.1.8.dist-info → python_codex-0.1.10.dist-info}/RECORD

@@ -3,10 +3,10 @@ pycodex/agent.py,sha256=s0FrF_XG2pHKryooS461Jr_acmQ_TKTp2JLGQNiny6w,11888
 pycodex/cli.py,sha256=MDd6aK9L3FkNSmwIiYQs322HPqL4L8kQl9laRn4brLM,31924
 pycodex/collaboration.py,sha256=yQ6pBD-R3ZWR4_FAYQFoS7KF0m4LLD42otXIbPqw2ys,641
 pycodex/compat.py,sha256=IO0X7AgcYhlHnYnpvBZ6leCh_UjoQzg5HLT5wYBNNIw,3155
-pycodex/context.py,sha256=R5tuMcNrX1F-Lh9ymsSbnfRbKLJ19TWrtQoZ3tWlHvM,24982
+pycodex/context.py,sha256=NuyiDJoUbhj4Xp6xExqs4cATGCKfYQ-YUx2BC7JXGeA,25306
 pycodex/doctor.py,sha256=De3M4hRBJq8ZeqsUJgHz0vitqrH18YugrEnz7oHhTdQ,10572
-pycodex/model.py,sha256=Mk9LZKmFcXG71I18-gs4dUWNn0GIM1rbMhFfKDut_3w,32790
-pycodex/portable.py,sha256=kZ5XVOMZq0l6xXsx3FY9C3DfB4Jra5Hw38qTMH0TEwg,15597
+pycodex/model.py,sha256=p4eiRx_PTu8YK36zWQgldMgjW4PXNsdHSr4k8XCYbVQ,32806
+pycodex/portable.py,sha256=gxl2E2h5uZJbasMEPPs-nyALFPIvX79T2ZYsu6vXZrg,15656
 pycodex/portable_server.py,sha256=6I3pQkWj3e_SFlDXY2mGdCPns1w_3PSxByBV9wv5epI,7331
 pycodex/protocol.py,sha256=LYDzJefu1tugqQzee4NuZzxhGAv3hXrNcnlw04CudAY,11106
 pycodex/runtime.py,sha256=gpDDxQKfp1Cqh1U0uslI3rCoXEN7XpJGuHlV-bsveM0,7983
@@ -28,7 +28,7 @@ pycodex/prompts/permissions/sandbox_mode/workspace_write.md,sha256=lVN-LwrBbHqlv
 pycodex/tools/__init__.py,sha256=aSLXrr_31KGQgDfRow5zVIc-2-KdXlHaCE6qUnE4HWI,1772
 pycodex/tools/agent_tool_schemas.py,sha256=r7pBICcx8fb0Rg6IzIg8-u3um2z11TogQ4yCzuiO-4o,2033
 pycodex/tools/apply_patch_tool.py,sha256=aFob-gzaCXlzPdCIvRXVKm1NrQqqhqe8CVkFVAhqiTc,13955
-pycodex/tools/base_tool.py,sha256=vlNw1PoQhXXb8evPjj7bqLOuOcw1ttKl-SrXOHJGVvs,6266
+pycodex/tools/base_tool.py,sha256=gClTp7i4KIwb8IGAAfhjqJbvEFKo-8QCe4rPXjgDLUM,6314
 pycodex/tools/close_agent_tool.py,sha256=nY3l_UOX6NyTgUqdXag3yRpdyQScV0g0Vv4HE3ElLwg,1597
 pycodex/tools/code_mode_manager.py,sha256=Wow42H_9IomUKUjjjU8rrAFAklhE-UlgxgrbgHRU_4M,19031
 pycodex/tools/exec_command_tool.py,sha256=l8GWlZKTvlWWAd_OPKsnnt3m0woMWXK8NkilmspnaQQ,3485
@@ -54,25 +54,26 @@ pycodex/tools/write_stdin_tool.py,sha256=nCuProkbeewfQ_yS8CgBajo--K3EmkXzJYh1D2Q
 pycodex/utils/__init__.py,sha256=jE63cZR1IBzs4ED86lwdYRLqV5FmPhNNDzIgC90mr6A,1216
 pycodex/utils/compactor.py,sha256=ZCzGc02xHmXq1rIjnG2gATKcFtt6r-OGsCIK0ypjnyI,6467
 pycodex/utils/debug.py,sha256=JeEB5JfzYfbdG0fXlrWFmXyR1ts86fKsI_97IqgF6R0,296
-pycodex/utils/dotenv.py,sha256=EDBXdn93ewmq9zhJki5_LsJJXe0wMIQJ6VfCE1r7voQ,1818
-pycodex/utils/get_env.py,sha256=jR8G0Xco57jX-71E1oHIcl3-Kz9Ltc0kzxj04DKzt80,7316
+pycodex/utils/dotenv.py,sha256=rGKmurHjm7GdP4giyjHBPpSPv2Oi45qBqDB6HG3CnfA,1866
+pycodex/utils/get_env.py,sha256=5fNhcNhujOakWV6AS66rGW3jEA68WGpuE4YVXJZFE6U,7427
 pycodex/utils/random_ids.py,sha256=zBphjVGc7OXk9ZNExAbxRi_bk7ipyLG491qTv7hi8jM,380
-pycodex/utils/session_persist.py,sha256=dUvo3Z1QBB4HJT1tLerDlLD3ZB25umB6FP6JORg9V40,16414
+pycodex/utils/session_persist.py,sha256=Ntu0jcb2cEZbXpKDe0LXD-OuxfjK0SzBV0lRi90-NAM,16496
 pycodex/utils/visualize.py,sha256=JURzq2AbV046bblE5fojcAe885Juda0LDxt_gqT2PUc,41006
 responses_server/__init__.py,sha256=3yPv_zeGT7P11tTnmj5kXktISLNsNW-02MUnnbiZcb0,394
 responses_server/__main__.py,sha256=9SRp-Yw7ShGxc6DhSIXcDLKgGEdAVm3oBZ59rBOPjT0,62
-responses_server/app.py,sha256=4SUG8xqKqmVKVY9i1f5WF9QrnmxGbD4mwiI6s13zQDE,7742
+responses_server/app.py,sha256=53SEwI2pUj_YL8zWvvYhQHiT1EUVVPzdRzOqXDvFMJ0,7770
 responses_server/config.py,sha256=leb3_uPrCyYdUIkyRyVPX4luGF88dQ62OkhRLPe7uxw,2718
-responses_server/messages_api.py,sha256=3GPMfs3ksQkhezLyWBjeW5zJ1e_MeHXVaq1lALIA7Mk,16815
-responses_server/payload_processors.py,sha256=gfOXqvVwlhCk-yjaDdGU4RKcpDdxIq2y6CmqUCggIjY,3444
-responses_server/server.py,sha256=isyzN-p-Ir8LLycN_dQfcanvie2ZqqSu52mOPz_wYD4,2095
+responses_server/messages_api.py,sha256=WgO6J1jz2pOJkI79rLXp-pS1yxtLARcwX8T6JX5Vkcc,16971
+responses_server/payload_processors.py,sha256=cbXGW8Xi-mliaWRg0_Af41X0vXV2W6R9VBzTE6DXfe4,3483
+responses_server/server.py,sha256=Ko-Cqz_kW-uve091itucMklsPhEei77v-YcTjtjEdqU,2286
 responses_server/session_store.py,sha256=ZD3cH2aEOkWaQsu5qTzcal2mThTSFQPAhAhPUN9srgI,1115
-responses_server/stream_router.py,sha256=OVwaDEsUaKVDNGF2vnqNZTo3WA9h3D3uzHeYY-QN9IU,34754
+responses_server/stream_router.py,sha256=PuW8_fo8c_R9kd_Gy_Z7nyij7-xW_lPcfFcnYI2PvQA,35434
+responses_server/trajectory_dump.py,sha256=XCwYaZZmlAxSsSXOfhk3zRvyfDpOHX5R8KzspScNFUM,3435
 responses_server/tools/__init__.py,sha256=ivsBSEy0SBUhY-Uea5v1XMLXShkwHdCVl0id-1FwdZg,150
 responses_server/tools/custom_adapter.py,sha256=LxO7ldydvR-GWachDz8GKC0Q8KGGFoFPbZxM0QvxuZ0,8350
 responses_server/tools/web_search.py,sha256=pm4ZUiHUfxc0bGY1kEvt-BCzDrZIyP24xzPUcga2ul0,8908
-python_codex-0.1.8.dist-info/METADATA,sha256=LAgQh2YnOizB3dQExftZEmRaIMNE_rL6L7rNTxLpzfA,15719
-python_codex-0.1.8.dist-info/WHEEL,sha256=KGYbc1zXlYddvwxnNty23BeaKzh7YuoSIvIMO4jEhvw,87
-python_codex-0.1.8.dist-info/entry_points.txt,sha256=sNUVakoVuTrzJH505ZgRTQxmtRRPUHV_EH0i6EbYTyM,45
-python_codex-0.1.8.dist-info/licenses/LICENSE,sha256=0X8ifk312hYAORM4hlzg8wVSEXYKNmiPgWlB1YIy2Nw,10926
-python_codex-0.1.8.dist-info/RECORD,,
+python_codex-0.1.10.dist-info/METADATA,sha256=K4ZxbwNtBlhsCjm0z3Ro15K1eA97rplNd8k2qGUJt7U,15720
+python_codex-0.1.10.dist-info/WHEEL,sha256=KGYbc1zXlYddvwxnNty23BeaKzh7YuoSIvIMO4jEhvw,87
+python_codex-0.1.10.dist-info/entry_points.txt,sha256=sNUVakoVuTrzJH505ZgRTQxmtRRPUHV_EH0i6EbYTyM,45
+python_codex-0.1.10.dist-info/licenses/LICENSE,sha256=0X8ifk312hYAORM4hlzg8wVSEXYKNmiPgWlB1YIy2Nw,10926
+python_codex-0.1.10.dist-info/RECORD,,

responses_server/app.py

@@ -177,6 +177,7 @@ class ManagedResponseServer:
             self._app,
             host=self._config.host,
             port=self._config.port,
+            loop="asyncio",
             log_level="error",
             access_log=False,
         )

responses_server/messages_api.py

@@ -66,6 +66,8 @@ def build_messages_request(
         "max_tokens": _resolve_max_tokens(outcomming_request),
         "stream": bool(outcomming_request.get("stream", True)),
     }
+    if isinstance(outcomming_request.get("return_token_ids"), bool):
+        payload["return_token_ids"] = bool(outcomming_request.get("return_token_ids"))
     if system_blocks:
         payload["system"] = system_blocks
 

responses_server/payload_processors.py

@@ -32,6 +32,7 @@ class OutgoingRequest(TypedDict):
     tools: 'Optional[typing.List[typing.Dict[str, object]]]'
     tool_choice: 'Optional[object]'
     parallel_tool_calls: 'Optional[bool]'
+    return_token_ids: 'Optional[bool]'
 
 
 PayloadPostProcessor = Callable[[OutgoingRequest], OutgoingRequest]

responses_server/server.py

@@ -3,6 +3,7 @@ from .config import CompatServerConfig
 from .payload_processors import post_process_outcomming_request
 from .session_store import SessionStore
 from .stream_router import StreamRouter
+from .trajectory_dump import TrajectoryDumpWriter
 import typing
 
 
@@ -16,6 +17,7 @@ class ResponseServer:
         self._config = config
         self._session_store = session_store or SessionStore()
         self._stream_router = stream_router or StreamRouter(config)
+        self._trajectory_dump = TrajectoryDumpWriter.from_env()
 
     @property
     def config(self) -> 'CompatServerConfig':
@@ -38,6 +40,9 @@ class ResponseServer:
         request_headers: 'typing.Dict[str, str]',
     ):
         outcomming_request = self._stream_router.build_outcomming_request(request_body)
+        if self._trajectory_dump is not None:
+            # vLLM surfaces prompt/decode token IDs only when this flag is set.
+            outcomming_request["return_token_ids"] = True
         outcomming_request = post_process_outcomming_request(
             outcomming_request,
             self._config.model_provider,
@@ -52,12 +57,9 @@ class ResponseServer:
             session_id=session_id,
             model=str(outcomming_request["model"]),
         )
-        incomming_stream = self._stream_router.open_outcomming_stream(
-            outcomming_request
-        )
         return self._stream_router.route_stream(
-            incomming_stream,
             stored_response,
             outcomming_request,
             custom_tool_names,
+            self._trajectory_dump,
         )

responses_server/stream_router.py

@@ -27,6 +27,7 @@ from .tools.web_search import (
     hydrate_tool_call_names,
     partition_tool_calls,
 )
+from .trajectory_dump import TrajectoryDumpWriter
 import typing
 
 
@@ -285,10 +286,10 @@ class StreamRouter:
 
     def route_stream(
         self,
-        incomming_stream,
         stored_response: 'StoredResponse',
         outcomming_request: 'typing.Dict[str, object]',
         custom_tool_names: 'typing.Union[typing.Set[str], None]' = None,
+        trajectory_dump: 'typing.Union[TrajectoryDumpWriter, None]' = None,
     ):
         yield (
             "response.created",
@@ -307,7 +308,10 @@ class StreamRouter:
         reasoning_parts: 'typing.List[str]' = []
         latest_usage: 'typing.Dict[str, object]' = {}
         current_request = json.loads(json.dumps(outcomming_request))
-        current_stream = incomming_stream
+        current_stream = self._open_tracked_outcomming_stream(
+            current_request,
+            trajectory_dump,
+        )
 
         while True:
             tool_calls: 'typing.Dict[int, typing.Dict[str, object]]' = {}
@@ -352,7 +356,10 @@ class StreamRouter:
                     )
                 except ValueError as exc:
                     raise OutcommingChatError(str(exc)) from exc
-                current_stream = self.open_outcomming_stream(current_request)
+                current_stream = self._open_tracked_outcomming_stream(
+                    current_request,
+                    trajectory_dump,
+                )
                 continue
 
             for item in self._build_output_items(
@@ -394,6 +401,16 @@ class StreamRouter:
             },
         )
 
+    def _open_tracked_outcomming_stream(
+        self,
+        outcomming_request: 'typing.Dict[str, object]',
+        trajectory_dump: 'typing.Union[TrajectoryDumpWriter, None]' = None,
+    ):
+        outcomming_stream = self.open_outcomming_stream(outcomming_request)
+        if trajectory_dump is None:
+            return outcomming_stream
+        return trajectory_dump.wrap_stream(outcomming_stream)
+
     def _responses_input_to_chat_messages(
         self,
         instructions: 'str',

responses_server/trajectory_dump.py

@@ -0,0 +1,105 @@
+import json
+import os
+import sys
+import threading
+import time
+import typing
+
+
+class TrajectoryDumpWriter:
+    ENV_VAR = "PYCODEX_DUMP"
+
+    def __init__(self, root_dir: 'str') -> 'None':
+        self._root_dir = os.path.abspath(root_dir)
+        self._dump_path = os.path.join(self._root_dir, "dump.jsonl")
+        self._lock = threading.Lock()
+        os.makedirs(self._root_dir, exist_ok=True)
+
+    @classmethod
+    def from_env(cls) -> 'typing.Union[TrajectoryDumpWriter, None]':
+        root_dir = str(os.environ.get(cls.ENV_VAR, "") or "").strip()
+        if not root_dir:
+            return None
+        return cls(root_dir)
+
+    def wrap_stream(self, outcomming_stream):
+        def iter_stream():
+            capture = _TrajectoryCapture(self, time.time())
+            try:
+                for chunk in outcomming_stream:
+                    capture.observe_chunk(chunk)
+                    yield chunk
+            finally:
+                capture.flush()
+
+        return iter_stream()
+
+    def _append_record(self, record: 'typing.Dict[str, object]') -> 'None':
+        serialized = json.dumps(record, ensure_ascii=False)
+        with self._lock:
+            os.makedirs(self._root_dir, exist_ok=True)
+            with open(self._dump_path, "a", encoding="utf-8") as handle:
+                handle.write(serialized)
+                handle.write("\n")
+
+
+class _TrajectoryCapture:
+    def __init__(
+        self,
+        writer: 'TrajectoryDumpWriter',
+        send_timestamp: 'float',
+    ) -> 'None':
+        self._writer = writer
+        self._send_timestamp = float(send_timestamp)
+        self._prefill_token_ids = None
+        self._decode_token_ids = []
+        self._closed = False
+
+    def observe_chunk(self, payload: 'object') -> 'None':
+        if not isinstance(payload, dict):
+            return
+        if self._prefill_token_ids is None and "prompt_token_ids" in payload:
+            normalized_prefill = _normalize_token_ids(payload.get("prompt_token_ids"))
+            if normalized_prefill is not None:
+                self._prefill_token_ids = normalized_prefill
+
+        choices = payload.get("choices") or []
+        if not isinstance(choices, list):
+            return
+        for raw_choice in choices:
+            if not isinstance(raw_choice, dict):
+                continue
+            normalized_decode = _normalize_token_ids(raw_choice.get("token_ids"))
+            if normalized_decode:
+                self._decode_token_ids.extend(normalized_decode)
+
+    def flush(self) -> 'None':
+        if self._closed:
+            return
+        self._closed = True
+        record = {
+            "tokens": {
+                "prefill": list(self._prefill_token_ids or []),
+                "decode": list(self._decode_token_ids),
+            },
+            "send_timestamp": self._send_timestamp,
+        }
+        try:
+            self._writer._append_record(record)
+        except Exception as exc:
+            print(
+                "responses_server: failed to append PYCODEX_DUMP trajectory: %s"
+                % exc,
+                file=sys.stderr,
+            )
+
+
+def _normalize_token_ids(raw_value: 'object') -> 'typing.Union[typing.List[int], None]':
+    if not isinstance(raw_value, list):
+        return None
+    token_ids = []
+    for value in raw_value:
+        if isinstance(value, bool) or not isinstance(value, int):
+            continue
+        token_ids.append(value)
+    return token_ids