diff --git a/integrations/codex/plugins/reme/hooks/auto_memory.py b/integrations/codex/plugins/reme/hooks/auto_memory.py index eb7eb3d0..1872c188 100644 --- a/integrations/codex/plugins/reme/hooks/auto_memory.py +++ b/integrations/codex/plugins/reme/hooks/auto_memory.py @@ -21,9 +21,12 @@ from reme_guidance import memory_guidance # noqa: E402 from reme_mcp import call # noqa: E402 from reme_state import digest, log_status, queue_root # noqa: E402 +sys.stdout.reconfigure(encoding="utf-8") + HOME_ENV = reme_config.HOME_ENV HOST = "codex" CAPTURE_WAIT_SECONDS = 10 +CAPTURE_SETTLE_SECONDS = 3.0 def clean_text(value: str) -> str: @@ -37,9 +40,12 @@ def transcript_messages(path: Path, *, root_only: bool = True): """Read current completed-item events, excluding reasoning and response-item duplicates.""" with path.open(encoding="utf-8") as handle: for index, line in enumerate(handle): - if not line.endswith("\n"): - break - record = json.loads(line) + try: + record = json.loads(line.rstrip("\r\n")) + except (json.JSONDecodeError, UnicodeError): + if not line.endswith("\n"): + return + continue if not isinstance(record, dict): continue payload = record.get("payload") @@ -83,19 +89,32 @@ def completed_turn(payload: dict, session: str, *, root_only: bool = True) -> li payload.get("transcript_path"), payload.get("turn_id"), ) - if not all(isinstance(value, str) and value.strip() for value in (final, transcript, turn_id)): + if not all(isinstance(value, str) and value.strip() for value in (transcript, turn_id)): + return [] + if final is not None and not isinstance(final, str): return [] path = Path(transcript).expanduser() if root_only and "subagents" in path.parts: return [] + try: + messages = list(transcript_messages(path, root_only=root_only)) + except FileNotFoundError: + return [] + final_text = clean_text(final) if isinstance(final, str) else "" + turn_messages = [message for message in messages if message["turn_id"] == turn_id] user, pair = None, [] - for message in transcript_messages(path, root_only=root_only): - if message["turn_id"] != turn_id: - continue + for message in turn_messages: if message["role"] == "user": user = message - elif user and message["text"] == clean_text(final): + elif user and final_text and message["text"] == final_text: pair = [user, message] + if not pair: + user = None + for message in turn_messages: + if message["role"] == "user": + user = message + elif user and message["role"] == "assistant": + pair = [user, message] return [ { "id": f"{HOST}-{digest(session + ':' + message['id'])}", @@ -108,6 +127,22 @@ def completed_turn(payload: dict, session: str, *, root_only: bool = True) -> li ] +def completed_turn_with_retry( + payload: dict, + session: str, + *, + root_only: bool = True, + timeout: float | None = None, +) -> list[dict]: + """Allow Codex a short window to finish writing the current transcript turn.""" + deadline = time.monotonic() + max(0, CAPTURE_SETTLE_SECONDS if timeout is None else timeout) + while True: + messages = completed_turn(payload, session, root_only=root_only) + if len(messages) == 2 or time.monotonic() >= deadline: + return messages + time.sleep(min(0.1, deadline - time.monotonic())) + + def session_queue(config: dict, payload: dict) -> Path | None: """Locate this agent's queue for the currently configured service without reading source text.""" native_id = payload.get("agent_id") or payload.get("session_id") @@ -117,12 +152,22 @@ def session_queue(config: dict, payload: dict) -> Path | None: return queue_root(config) / session +def delivery_session_id(payload: dict) -> str | None: + """Use the stable native id for the session store while keeping queue paths private.""" + native_id = payload.get("agent_id") or payload.get("session_id") + if not isinstance(native_id, str) or not native_id: + return None + safe = re.sub(r'[<>:"/\\|?*]', "-", native_id)[:120] + return f"{HOST}-{safe}" + + def capture(config: dict, payload: dict) -> Path | None: """Persist a completed turn before any network request.""" root = session_queue(config, payload) - if root is None: + session_id = delivery_session_id(payload) + if root is None or session_id is None: return None - messages = completed_turn(payload, root.name, root_only=config["rootAgentsOnly"]) + messages = completed_turn_with_retry(payload, root.name, root_only=config["rootAgentsOnly"]) if len(messages) != 2: log_status("capture_skipped", config=config) return None @@ -132,7 +177,7 @@ def capture(config: dict, payload: dict) -> Path | None: day = instant.astimezone(ZoneInfo(config["timezone"])).date().isoformat() write_json( root / f"{key}.json", - {"session_id": root.name, "turn_id": payload["turn_id"], "messages": messages, "date": day}, + {"session_id": session_id, "turn_id": payload.get("turn_id"), "messages": messages, "date": day}, ) log_status("memory_queued", config=config, turns=1) return root @@ -225,7 +270,7 @@ def recall(config: dict, payload: dict) -> dict: f"{text}\n" ) except (OSError, ValueError, RuntimeError) as exc: - log_status("recall_failed", exc, config=config) + log_status("recall_failed", exc.__cause__ or exc, config=config) else: log_status("recall_disabled", config=config) return {"hookSpecificOutput": {"hookEventName": "UserPromptSubmit", "additionalContext": context}} diff --git a/tests/unit/test_codex_capture_recovery.py b/tests/unit/test_codex_capture_recovery.py new file mode 100644 index 00000000..dac7b27d --- /dev/null +++ b/tests/unit/test_codex_capture_recovery.py @@ -0,0 +1,116 @@ +"""Recovery behavior for Codex Stop hooks when transcript writes settle late.""" + +# pylint: disable=redefined-outer-name + +import importlib.util +import json +from pathlib import Path + +import pytest + +PLUGIN = Path(__file__).resolve().parents[2] / "integrations/codex/plugins/reme" + + +@pytest.fixture +def codex_hook(tmp_path, monkeypatch): + """Load the Codex hook against an isolated host configuration.""" + spec = importlib.util.spec_from_file_location("codex_capture_recovery", PLUGIN / "hooks/auto_memory.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + monkeypatch.setenv(module.HOME_ENV, str(tmp_path / "host")) + monkeypatch.setattr(module, "CAPTURE_SETTLE_SECONDS", 0) + return module + + +def completed(turn_id, item_type, text): + """Build one completed Codex transcript item.""" + return { + "type": "event_msg", + "timestamp": "2026-10-26T01:00:00Z", + "payload": { + "type": "item_completed", + "turn_id": turn_id, + "item": { + "type": item_type, + "content": [{"type": "text" if item_type == "UserMessage" else "Text", "text": text}], + }, + }, + } + + +def test_accepts_unterminated_final_transcript_line(codex_hook, tmp_path): + """Accept a final row that is complete JSON but has no trailing newline.""" + path = tmp_path / "rollout.jsonl" + rows = [completed("turn-one", "UserMessage", "question"), completed("turn-one", "AgentMessage", "answer")] + path.write_text("\n".join(json.dumps(row) for row in rows), encoding="utf-8") + payload = {"transcript_path": str(path), "last_assistant_message": "answer", "turn_id": "turn-one"} + pair = codex_hook.completed_turn(payload, "session") + assert [message["content"][0]["text"] for message in pair] == ["question", "answer"] + + +def test_uses_last_same_turn_assistant_when_stop_text_differs(codex_hook, tmp_path): + """Fall back to the last same-turn agent message when Stop text differs.""" + path = tmp_path / "rollout.jsonl" + rows = [ + completed("turn-one", "UserMessage", "question"), + completed("turn-one", "AgentMessage", "progress"), + completed("turn-one", "AgentMessage", "final answer"), + ] + path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + payload = { + "transcript_path": str(path), + "last_assistant_message": "rendered final answer", + "turn_id": "turn-one", + } + pair = codex_hook.completed_turn(payload, "session") + assert [message["content"][0]["text"] for message in pair] == ["question", "final answer"] + payload["last_assistant_message"] = None + pair = codex_hook.completed_turn(payload, "session") + assert [message["content"][0]["text"] for message in pair] == ["question", "final answer"] + + +def test_does_not_cross_turns_when_turn_id_is_unknown(codex_hook, tmp_path): + """Never capture messages from another turn.""" + path = tmp_path / "rollout.jsonl" + rows = [completed("turn-one", "UserMessage", "question"), completed("turn-one", "AgentMessage", "answer")] + path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + payload = {"transcript_path": str(path), "last_assistant_message": "answer", "turn_id": "unknown"} + assert codex_hook.completed_turn(payload, "session") == [] + + +def test_capture_retries_turn_materialization(codex_hook, monkeypatch): + """Retry until the current turn is fully materialized.""" + pair = [ + { + "id": "codex-1", + "role": "user", + "content": [{"type": "text", "text": "question"}], + "created_at": "2026-10-26T01:00:00+00:00", + }, + { + "id": "codex-2", + "role": "assistant", + "content": [{"type": "text", "text": "answer"}], + "created_at": "2026-10-26T01:00:01+00:00", + }, + ] + responses = iter([[], pair]) + monkeypatch.setattr(codex_hook, "completed_turn", lambda *_args, **_kwargs: next(responses)) + assert codex_hook.completed_turn_with_retry({}, "session", timeout=0.2) == pair + + +def test_capture_delivers_with_native_session_id(codex_hook, tmp_path): + """Deliver with the native session ID rather than the queue hash.""" + path = tmp_path / "rollout.jsonl" + rows = [completed("turn-one", "UserMessage", "question"), completed("turn-one", "AgentMessage", "answer")] + path.write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8") + payload = { + "hook_event_name": "Stop", + "session_id": "01a1766c-4b23-7df3-914c-32d0020e7285", + "transcript_path": str(path), + "last_assistant_message": "answer", + "turn_id": "turn-one", + } + root = codex_hook.capture(codex_hook.load_config(), payload) + queued = json.loads(next(root.glob("*.json")).read_text(encoding="utf-8")) + assert queued["session_id"] == "codex-01a1766c-4b23-7df3-914c-32d0020e7285"