From e0763d1a3d09c0232da7876b9542ffa12aede0a0 Mon Sep 17 00:00:00 2001 From: LoopX Agent <337587101+loopx-agent@users.noreply.github.com> Date: Thu, 5 Nov 2026 13:39:44 +0800 Subject: [PATCH] Bind prepared context delivery to the lifecycle event session Signed-off-by: LoopX Agent <337587101+loopx-agent@users.noreply.github.com> --- loopx/cli_commands/quota_capture.py | 13 ++++--- .../turn_driver/codex_context.py | 8 ++-- .../test_quota_execution_preparation.py | 37 ++++++++++++++----- 3 files changed, 39 insertions(+), 19 deletions(-) diff --git a/loopx/cli_commands/quota_capture.py b/loopx/cli_commands/quota_capture.py index 9c598223b0..f385f287c8 100644 --- a/loopx/cli_commands/quota_capture.py +++ b/loopx/cli_commands/quota_capture.py @@ -457,6 +457,7 @@ def _native_context_read_bytes(reference: dict[str, Any], raw: bytes, *, def native_context_capture_view(projected: dict[str, Any], directory: Path, *, context_directory: str | None = None, + hook_session: str | None = None, captured_recovery: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, max_page_output_bytes: int | None = None, @@ -545,7 +546,7 @@ def view(kind: str, ref: dict[str, Any], body: bytes) -> dict[str, Any]: from ..control_plane.turn_driver.codex_context import register_native_context_view register_native_context_view(context_directory, - {**result, "ancestry": ancestry}) + {**result, "ancestry": ancestry}, hook_session=hook_session) return result @@ -613,7 +614,7 @@ def progressive_context_capture_view(payload: dict[str, Any], args: argparse.Nam return readable if len(_native_output_bytes(readable)) < len(_native_output_bytes(required)) else required -def project_saved_context_capture(args: argparse.Namespace) -> dict[str, Any]: +def project_saved_context_capture(args: argparse.Namespace, *, hook_session: str | None = None) -> dict[str, Any]: """Present an immutable already-decided quota capture, without authority IO. Publication is delayed until this native tool runs, so the lifecycle owner @@ -677,7 +678,7 @@ def project_saved_context_capture(args: argparse.Namespace) -> dict[str, Any]: page_bytes=page_bytes, output_bytes=output_bytes, reader_argv=reader_argv) reference = dict(full_reference) directory = codex_context_directory(directory=args.codex_context_dir, - root=args.codex_context_root, create=True) + root=args.codex_context_root, session=hook_session, create=True) if not directory: raise ValueError("native context is unavailable") # The session already owns a private recovery location. Callers need @@ -700,7 +701,7 @@ def project_saved_context_capture(args: argparse.Namespace) -> dict[str, Any]: goal_id=args.goal_id, agent_id=args.agent_id, surface="quota_packet", capture_id=hashlib.sha256(raw).hexdigest(), continuity_generation=getattr(args, "_context_delivery_generation", None), - progressive_enabled=activation["progressive_enabled"]) + progressive_enabled=activation["progressive_enabled"], hook_session=hook_session) native_prepared = True if projected.get("schema_version") == "context_delivery_v1": # Native publication owns the complete required view. Keep its @@ -709,7 +710,7 @@ def project_saved_context_capture(args: argparse.Namespace) -> dict[str, Any]: full_reference = _capture_read_reference(projected["recovery"], page_bytes=page_bytes, output_bytes=output_bytes, reader_argv=reader_argv) reference = dict(full_reference) - native_view = native_context_capture_view(projected, cache, context_directory=directory, + native_view = native_context_capture_view(projected, cache, context_directory=directory, hook_session=hook_session, metadata={"recovery": full_reference, "capture_sha256": args.capture_sha256, "context_delivery": activation}, max_page_output_bytes=getattr(args, "max_page_output_bytes", None), reader_argv=reader_argv) @@ -1184,7 +1185,7 @@ def handle_native_context_command(args: argparse.Namespace, *, print_payload: Ca "capture_sha256": hashlib.sha256(source.read_bytes()).hexdigest(), "goal_id": request["goal_id"], "agent_id": request["agent_id"], "turn_instance_id": request["turn_instance_id"], - })) + }), hook_session=event["session_id"]) existing = payload.get("hookSpecificOutput") existing_context = ( str(existing.get("additionalContext") or "") diff --git a/loopx/control_plane/turn_driver/codex_context.py b/loopx/control_plane/turn_driver/codex_context.py index 39d80fd8cd..99a3f9c84b 100644 --- a/loopx/control_plane/turn_driver/codex_context.py +++ b/loopx/control_plane/turn_driver/codex_context.py @@ -578,7 +578,7 @@ def _native_context_request(state: dict[str, Any], scope: dict[str, Any]) -> dic return request -def register_native_context_view(directory: str | None, view: dict[str, Any]) -> None: +def register_native_context_view(directory: str | None, view: dict[str, Any], *, hook_session: str | None = None) -> None: """Bind the existing pending delivery to its actual paged wire, under lock.""" if directory is None or not isinstance(view.get("target"), dict): return @@ -586,7 +586,7 @@ def register_native_context_view(directory: str | None, view: dict[str, Any]) -> state = _state(root) pending = state.get("pending") if (not isinstance(pending, dict) or pending.get("target") != view["target"] - or state.get("session_id") != os.environ.get("CODEX_THREAD_ID")): + or state.get("session_id") != (hook_session or os.environ.get("CODEX_THREAD_ID"))): raise ValueError("native page view does not match the current pending capture") pending.update(path=view["path"], sha256=view["sha256"], total_bytes=view["total_bytes"], ancestry=view.get("ancestry", 0), @@ -597,7 +597,7 @@ def register_native_context_view(directory: str | None, view: dict[str, Any]) -> def project_codex_context(payload: dict[str, Any], *, directory: str | None, capture_directory: Path, goal_id: str, agent_id: str, surface: str, capture_id: str, continuity_generation: str | None = None, - progressive_enabled: bool = False) -> dict[str, Any]: + progressive_enabled: bool = False, hook_session: str | None = None) -> dict[str, Any]: """Commit current recovery before optional projection or candidate updates. Serialization prevents invalidation from moving the epoch during preparation. @@ -620,7 +620,7 @@ def project_codex_context(payload: dict[str, Any], *, directory: str | None, cap try: state = _state(root) session, epoch = state.get("session_id"), state.get("context_epoch") - native_session = os.environ.get("CODEX_THREAD_ID") + native_session = hook_session or os.environ.get("CODEX_THREAD_ID") if (not isinstance(session, str) or not session or (native_session and session != native_session)): return payload # Do not replace a different native session's recovery. diff --git a/tests/control_plane/test_quota_execution_preparation.py b/tests/control_plane/test_quota_execution_preparation.py index 214ec8c6f8..95fcf3cdf6 100644 --- a/tests/control_plane/test_quota_execution_preparation.py +++ b/tests/control_plane/test_quota_execution_preparation.py @@ -175,18 +175,28 @@ def test_projected_claim_prepares_the_same_todo_for_a_fresh_guard( assert after["normal_delivery_allowed"] is True -@pytest.mark.parametrize("provider", ["file", "sqlite"]) -@pytest.mark.parametrize("scenario", ["codex_app_heartbeat", "codex_app_ssh_goal"]) +@pytest.mark.parametrize("provider,scenario,location,ambient", [ + (provider, scenario, "directory", "preparation-session") + for provider in ("file", "sqlite") + for scenario in ("codex_app_heartbeat", "codex_app_ssh_goal") +] + [ + ("sqlite", "codex_app_heartbeat", "root", None), + ("sqlite", "codex_app_heartbeat", "root", "other-session"), +]) def test_native_hook_returns_readable_ready_source_and_original_claim( - tmp_path, monkeypatch, provider, scenario, + tmp_path, monkeypatch, provider, scenario, location, ambient, ): from loopx.control_plane.turn_driver import codex_context tmp_path = tmp_path.resolve() monkeypatch.chdir(tmp_path) monkeypatch.setenv("CODEX_THREAD_ID", "preparation-session") - native = tmp_path / "native" - native.mkdir(mode=0o700) + root = tmp_path / "native" + root.mkdir(mode=0o700) + native = root + if location == "root": + native = Path(codex_context.codex_context_directory(directory=None, root=str(root), + session="preparation-session", create=True)) event = {"session_id": "preparation-session", "cwd": str(tmp_path), "hook_event_name": "SessionStart", "source": "startup"} assert codex_context.codex_context_hook(str(native), event) == {} @@ -204,15 +214,21 @@ def test_native_hook_returns_readable_ready_source_and_original_claim( goal_id=GOAL, agent_id=AGENT, surface="quota_packet", capture_id="before-claim") hook = [sys.executable, "-m", "loopx.entrypoint", "--registry", str(tmp_path / "registry.json"), "--runtime-root", str(tmp_path / "runtime"), "--format", "json", "quota", "context-hook", - "--context-delivery", "on", "--runtime-profile", scenario, "--codex-context-dir", str(native)] + "--context-delivery", "on", "--runtime-profile", scenario, + "--codex-context-root" if location == "root" else "--codex-context-dir", str(root)] + hook_env = os.environ.copy() + if ambient is None: + hook_env.pop("CODEX_THREAD_ID", None) + else: + hook_env["CODEX_THREAD_ID"] = ambient hook_event = {**event, "hook_event_name": "PostToolUse"} log = tmp_path / "runtime" / "goals" / GOAL / "rollout-event-log.jsonl" unchanged = log.read_bytes() disabled = subprocess.run(["off" if item == "on" else item for item in hook], - input=json.dumps(hook_event), cwd=tmp_path, capture_output=True, text=True, timeout=30) + input=json.dumps(hook_event), env=hook_env, cwd=tmp_path, capture_output=True, text=True, timeout=30) assert disabled.returncode == 0 and json.loads(disabled.stdout) == {} assert log.read_bytes() == unchanged - completed = subprocess.run(hook, input=json.dumps(hook_event), cwd=tmp_path, + completed = subprocess.run(hook, input=json.dumps(hook_event), env=hook_env, cwd=tmp_path, capture_output=True, text=True, timeout=30) assert completed.returncode == 0, completed.stderr additional = json.loads(completed.stdout)["hookSpecificOutput"]["additionalContext"] @@ -223,6 +239,9 @@ def test_native_hook_returns_readable_ready_source_and_original_claim( claim_reference = receipt["claim_operation_id"] assert claim_reference.startswith("todo-claim:") ready = receipt["context"] + assert ready["ok"], ready.get("reason") + if location == "root": + assert list(root.iterdir()) == [native] read = shlex.split(ready["recovery"]["command"]) pointers = ["/execution_preparation", "/selected_todo", "/heartbeat_receipt", "/interaction_contract/cli_channel", "/interaction_contract/agent_channel/work_context", "/interaction_contract/agent_channel/required_reads"] @@ -259,7 +278,7 @@ def test_native_hook_returns_readable_ready_source_and_original_claim( registry_path=tmp_path / "registry.json", runtime_root=tmp_path / "runtime", cwd=tmp_path) assert code == 0 and progress["ok"], progress events = log.read_bytes() - repeat = subprocess.run(hook, input=json.dumps(hook_event), cwd=tmp_path, + repeat = subprocess.run(hook, input=json.dumps(hook_event), env=hook_env, cwd=tmp_path, capture_output=True, text=True, timeout=30) assert repeat.returncode == 0 and "host_execution_prepared" not in repeat.stdout assert log.read_bytes() == events