"""The run-scoped work branch (p2hs/codegen/): naming, state, and the resume decision.""" from __future__ import annotations from datetime import UTC, datetime, timedelta from types import SimpleNamespace from unittest.mock import AsyncMock from uuid import uuid4 import pytest from src.domain.codegen.work_branch import ( AGENT_OUTPUT, CAPABILITIES_PROVISIONED, PARTIAL_WORK, PROMOTED, apply_report, bounded_correction, build_correction_appendix, is_expired, is_work_branch, new_state, persist_state, resume_plan, valid_state, work_branch_name, ) BASE = "a" * 40 CKPT = "b" * 40 CKPT2 = "c" * 40 def state(run_id=None, **overrides): return {**new_state(run_id or uuid4(), "preview", BASE), **overrides} # ---- naming / trust ----------------------------------------------------------------------- def test_branch_name_is_the_run_scoped_p2hs_codegen_ref(): run_id = uuid4() assert work_branch_name(run_id) == f"p2hs/codegen/{run_id}" assert is_work_branch(work_branch_name(run_id)) @pytest.mark.parametrize( "name", ["preview", "main", "p2hs/codegen/", "p2hs/codegen/../main", "p2hs/codegen/not-a-uuid", "refs/heads/p2hs/codegen/" + str(uuid4()), "p2hs/wip/" + str(uuid4()), "", None], ) def test_only_a_real_run_branch_is_a_work_branch(name): assert not is_work_branch(name) def test_a_state_is_trusted_only_for_its_own_run_and_a_real_base(): run_id = uuid4() good = state(run_id) assert valid_state(good, run_id) assert not valid_state(good, uuid4()) # another run's branch assert not valid_state({**good, "base_sha": "nope"}) assert not valid_state({**good, "branch": "preview"}) assert not valid_state(None) and not valid_state("x") # ---- checkpoints only move forward ---------------------------------------------------------- def test_reports_record_the_checkpoint_and_clear_failure_on_success(): failed = apply_report(state(), status="failed", checkpoint=AGENT_OUTPUT, checkpoint_sha=CKPT, failure_class="internal_error", retryable=True, message="boom") assert failed["last_checkpoint"] == AGENT_OUTPUT and failed["last_checkpoint_sha"] == CKPT assert failed["retryable"] is True and failed["failure_class"] == "internal_error" done = apply_report(failed, status="completed", checkpoint=PROMOTED, checkpoint_sha=CKPT2) assert done["failure_class"] is None and done["retryable"] is False and done["correction"] is None assert done["last_checkpoint"] == PROMOTED def test_a_stale_or_repeated_callback_cannot_rewind_the_checkpoint(): ahead = apply_report(state(), status="failed", checkpoint=CAPABILITIES_PROVISIONED, checkpoint_sha=CKPT2) behind = apply_report(ahead, status="failed", checkpoint=AGENT_OUTPUT, checkpoint_sha=CKPT) assert behind["last_checkpoint"] == CAPABILITIES_PROVISIONED assert behind["last_checkpoint_sha"] == CKPT2 def test_a_bogus_checkpoint_or_sha_is_ignored(): for checkpoint, sha in ((AGENT_OUTPUT, "not-a-sha"), ("made_up", CKPT), (None, CKPT)): new = apply_report(state(), status="failed", checkpoint=checkpoint, checkpoint_sha=sha) assert new["last_checkpoint"] is None # ---- resume decision ------------------------------------------------------------------------- def failed_state(failure_class, **extra): return apply_report( state(), status="failed", checkpoint=AGENT_OUTPUT, checkpoint_sha=CKPT, failure_class=failure_class, retryable=True, **extra, ) @pytest.mark.parametrize( "failure_class", ["internal_error", "cloud_unavailable", "cloud_rejected", "capability_gap", "preview_diverged"], ) def test_a_failure_after_the_agent_finished_resumes_without_re_running_it(failure_class): plan = resume_plan(failed_state(failure_class), max_corrections=2) assert plan == {"skip_agent": True, "last_checkpoint": AGENT_OUTPUT} def test_an_invalid_manifest_sends_the_agent_back_with_the_errors_up_to_a_limit(): correction = {"entries": [{"name": "contact_submission", "errors": [{"path": "fields.0.public", "code": "unknown_key", "message": "not recognized"}]}]} s = failed_state("invalid_requirement", correction=correction) plan = resume_plan(s, max_corrections=2) assert plan["skip_agent"] is False and plan["correction"]["entries"][0]["name"] == "contact_submission" assert resume_plan({**s, "correction_attempts": 2}, max_corrections=2) is None # give up @pytest.mark.parametrize("failure_class", ["internal_error", "invalid_requirement", "preview_diverged"]) def test_work_saved_before_the_agent_finished_re_runs_the_agent_on_top_of_it(failure_class): # The agent never completed, so its output cannot be validated or promoted as it stands: # a retry must run it again (on the saved branch), never skip it, whatever the failure class. saved = apply_report( state(), status="failed", checkpoint=PARTIAL_WORK, checkpoint_sha=CKPT, failure_class=failure_class, retryable=True, correction={"entries": [{"name": "x", "errors": []}]}, ) assert saved["last_checkpoint"] == PARTIAL_WORK assert resume_plan(saved, max_corrections=2) is None def test_partial_work_ranks_below_every_real_checkpoint(): finished = apply_report(state(), status="failed", checkpoint=AGENT_OUTPUT, checkpoint_sha=CKPT2) late = apply_report(finished, status="failed", checkpoint=PARTIAL_WORK, checkpoint_sha=CKPT) assert late["last_checkpoint"] == AGENT_OUTPUT and late["last_checkpoint_sha"] == CKPT2 early = apply_report(state(), status="failed", checkpoint=PARTIAL_WORK, checkpoint_sha=CKPT) advanced = apply_report(early, status="failed", checkpoint=AGENT_OUTPUT, checkpoint_sha=CKPT2) assert advanced["last_checkpoint"] == AGENT_OUTPUT def test_the_callback_schema_accepts_the_partial_work_checkpoint(): from src.domain.codegen.schemas import CodegenActionsCompleteBody body = CodegenActionsCompleteBody( status="failed", checkpoint="partial_work", checkpoint_sha=CKPT, error_message="boom" ) assert body.checkpoint == "partial_work" def test_nothing_to_resume_without_a_checkpoint_or_after_success(): assert resume_plan(state(), max_corrections=2) is None assert resume_plan(apply_report(state(), status="completed", checkpoint=PROMOTED, checkpoint_sha=CKPT), max_corrections=2) is None assert resume_plan(None, max_corrections=2) is None # ---- correction payload ----------------------------------------------------------------------- def test_the_correction_is_bounded_and_drops_anything_unstructured(): huge = {"entries": [{"name": "x" * 500, "errors": [{"path": "p" * 500, "code": "c", "message": "m" * 900, "secret": "leak"}] * 50, "renamed": [{"path": "name", "from": "A", "to": "a", "extra": 1}]}] * 50, "transcript": "do not keep"} bounded = bounded_correction(huge) assert len(bounded["entries"]) == 20 and len(bounded["entries"][0]["errors"]) == 20 assert len(bounded["entries"][0]["name"]) == 100 assert len(bounded["entries"][0]["errors"][0]["message"]) == 300 assert "secret" not in bounded["entries"][0]["errors"][0] and "transcript" not in bounded assert bounded_correction({"entries": []}) is None and bounded_correction("x") is None def test_the_correction_prompt_lists_errors_and_tells_the_agent_to_follow_renames(): text = build_correction_appendix( {"entries": [{"name": "contact_submission", "errors": [{"path": "fields.0.public", "code": "unknown_key", "message": "'public' is not a recognized key."}], "renamed": [{"path": "name", "from": "Contact Submission", "to": "contact_submission"}]}]}, attempt=2, ) assert "attempt 2" in text and "already in this repository" in text assert "fields.0.public" in text and "unknown_key" in text assert "`Contact Submission` to `contact_submission`" in text assert "Do not regenerate" in text and "never worked around" in text # ---- persistence / retention -------------------------------------------------------------------- @pytest.mark.asyncio async def test_the_state_is_stored_on_the_plan_context_and_other_context_survives(): run = SimpleNamespace(plan={"context": {"keep": "me"}, "system_prompt": "s"}) committed = [] class Session: async def __aenter__(self): return SimpleNamespace(commit=AsyncMock(side_effect=lambda: committed.append(1))) async def __aexit__(self, *args): return False from unittest.mock import patch with patch( "src.domain.codegen.repository.CodegenRunRepository", lambda db: SimpleNamespace(get_by_id=AsyncMock(return_value=run)), ): ok = await persist_state(lambda: Session(), uuid4(), state()) assert ok and committed assert run.plan["context"]["keep"] == "me" and run.plan["system_prompt"] == "s" assert run.plan["context"]["work_branch"]["base_branch"] == "preview" @pytest.mark.asyncio async def test_a_persistence_failure_never_raises(): def broken(): raise RuntimeError("db down") assert await persist_state(broken, uuid4(), state()) is False def test_only_finished_branches_expire_and_only_after_the_retention_period(): now = datetime.now(UTC) old = (now - timedelta(days=15)).isoformat() recent = (now - timedelta(days=2)).isoformat() assert is_expired(state(status="failed", updated_at=old), now, 14) assert not is_expired(state(status="failed", updated_at=recent), now, 14) assert not is_expired(state(status="active", updated_at=old), now, 14) # a live run assert not is_expired(state(status="failed", updated_at="garbage"), now, 14) assert not is_expired(None, now, 14)