Merge 00ff56e194 into c6049c3fc1

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")
@@ -76,26 +82,45 @@ def transcript_messages(path: Path, *, root_only: bool = True):
             }
 
 
-def completed_turn(payload: dict, session: str, *, root_only: bool = True) -> list[dict]:
-    """Match this Stop's final reply to its native user turn, even after another turn starts."""
+def _completed_turn(
+    payload: dict,
+    session: str,
+    *,
+    root_only: bool,
+    allow_fallback: bool,
+) -> list[dict]:
+    """Build a turn, optionally falling back after the caller's bounded retry window."""
     final, transcript, turn_id = (
         payload.get("last_assistant_message"),
         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 and allow_fallback and final_text:
+        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 +133,30 @@ def completed_turn(payload: dict, session: str, *, root_only: bool = True) -> li
     ]
 
 
+def completed_turn(payload: dict, session: str, *, root_only: bool = True) -> list[dict]:
+    """Match this Stop's final reply to its native user turn without accepting an early fallback."""
+    return _completed_turn(payload, session, root_only=root_only, allow_fallback=False)
+
+
+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:
+            return messages
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            return _completed_turn(payload, session, root_only=root_only, allow_fallback=True)
+        time.sleep(min(0.1, remaining))
+
+
 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")
@@ -122,7 +171,7 @@ def capture(config: dict, payload: dict) -> Path | None:
     root = session_queue(config, payload)
     if root 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 +181,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": root.name, "turn_id": payload.get("turn_id"), "messages": messages, "date": day},
         )
         log_status("memory_queued", config=config, turns=1)
     return root
@@ -225,7 +274,7 @@ def recall(config: dict, payload: dict) -> dict:
                     f"{text}\n</reme-context>"
                 )
         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}}

tests/unit/test_codex_capture_recovery.py

@@ -0,0 +1,137 @@
+"""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 only after the retry window when completed 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",
+    }
+    assert codex_hook.completed_turn(payload, "session") == []
+    pair = codex_hook.completed_turn_with_retry(payload, "session", timeout=0)
+    assert [message["content"][0]["text"] for message in pair] == ["question", "final answer"]
+    payload["last_assistant_message"] = None
+    assert codex_hook.completed_turn_with_retry(payload, "session", timeout=0) == []
+
+
+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_waits_for_final_item_instead_of_accepting_phase_less_progress(codex_hook, tmp_path, monkeypatch):
+    """Do not accept a phase-less intermediate item while the final reply may still be pending."""
+    path = tmp_path / "rollout.jsonl"
+    question = completed("turn-one", "UserMessage", "question")
+    progress = completed("turn-one", "AgentMessage", "progress")
+    final = completed("turn-one", "AgentMessage", "final answer")
+    path.write_text("".join(json.dumps(row) + "\n" for row in (question, progress)), encoding="utf-8")
+    payload = {
+        "transcript_path": str(path),
+        "last_assistant_message": "final answer",
+        "turn_id": "turn-one",
+    }
+
+    def append_final(_seconds):
+        path.write_text("".join(json.dumps(row) + "\n" for row in (question, progress, final)), encoding="utf-8")
+
+    monkeypatch.setattr(codex_hook.time, "sleep", append_final)
+    pair = codex_hook.completed_turn_with_retry(payload, "session", timeout=0.2)
+    assert [message["content"][0]["text"] for message in pair] == ["question", "final answer"]
+
+
+def test_capture_preserves_queue_session_id(codex_hook, tmp_path):
+    """Keep the queue hash as the delivered session identity."""
+    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"] == root.name