Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 7 additions & 6 deletions loopx/cli_commands/quota_capture.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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


Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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 "")
Expand Down
8 changes: 4 additions & 4 deletions loopx/control_plane/turn_driver/codex_context.py
Original file line number Diff line number Diff line change
Expand Up @@ -578,15 +578,15 @@ 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
with _locked(directory) as root:
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),
Expand All @@ -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.
Expand All @@ -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.
Expand Down
37 changes: 28 additions & 9 deletions tests/control_plane/test_quota_execution_preparation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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) == {}
Expand All @@ -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"]
Expand All @@ -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"]
Expand Down Expand Up @@ -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
Expand Down
Loading