from __future__ import annotations import json import zipfile from contextlib import contextmanager from pathlib import Path import pytest from django.test import Client from django.urls import reverse from control_plane.authoring.epub import write_epub from control_plane.authoring.models import ( CanonSnapshot, Chapter, ChapterRevision, ChapterStateDocument, EditorialFinding, GenerationContextSnapshot, OutlineVersion, StateChange, StateChangeStatus, StoryBibleVersion, StoryEntity, StoryProject, ) from control_plane.authoring.runner import StoryWorkflowRunner from control_plane.authoring.services import ( DjangoStoryWorkflowServices, apply_exact_edits, compact_chapter_plan, compact_scene_contract, deterministic_temporal_findings, scene_draft_packet, ) from control_plane.authoring.state_management import ( apply_state_changes, build_contract_requirements, evidence_is_present, ) from control_plane.authoring.streaming import DraftResult, ResumableDraftWriter from control_plane.authoring.workflow import build_story_workflow from control_plane.projects.models import Project from graph.models import GraphApproval, GraphApprovalStatus, GraphRunStatus from model_router.router import ( ModelChunk, ModelRequestContract, ModelResponseContract, ModelRouter, ) def story_revision() -> ChapterRevision: project = Project.objects.create(name="Story Test", project_type="STORY", goal="Write") story = StoryProject.objects.create(project=project, title="Test Book", slug="test-book") bible = StoryBibleVersion.objects.create(story=story, version=1, content="Canon") outline = OutlineVersion.objects.create( story=story, version=1, content={"chapters": [{"number": 1, "title": "Opening", "beats": ["Begin"]}]}, ) chapter = Chapter.objects.create(story=story, number=1, title="Opening") return ChapterRevision.objects.create( chapter=chapter, revision=1, story_bible=bible, outline=outline ) class FakeStoryServices: def __init__(self) -> None: self.review_calls: list[str] = [] self.audit_calls = 0 self.plan_calls = 0 def build_context(self, state): return {"context_snapshot_id": "context-1"} def plan_chapter(self, state): self.plan_calls += 1 return {"scene_plan": {"scenes": [{"number": 1}]}, "human_notes": ""} def draft_chapter(self, state): return {"revision_id": state["revision_id"]} def extract_continuity(self, state): return {} def judge_state_contract(self, state): self.audit_calls += 1 return {"state_judge_status": "pass", "editorial_finding_ids": []} def decide_patch(self, state): return {"patch_decision": "human_review", "patch_finding_ids": []} def apply_automatic_patch(self, state): raise AssertionError("patch should not run") def verify_patch(self, state): raise AssertionError("verification should not run") def state_approval_payload(self, state): return {} def review_chapter(self, state, review_kind): self.review_calls.append(review_kind) return [f"finding-{review_kind}"] def judge_chapter(self, state): return "human_review" def revise_chapter(self, state): return { "revision_id": state["revision_id"], "revision_attempt": int(state.get("revision_attempt") or 0) + 1, } def ensure_approval(self, state, gate, payload): approval, _ = GraphApproval.objects.get_or_create( graph_run_id=state["graph_run_id"], reason=f"{gate}:{state['revision_id']}", status=GraphApprovalStatus.PENDING, defaults={"payload": payload}, ) return approval def decide_approval(self, approval_id, decision): approval = GraphApproval.objects.get(id=approval_id) approval.status = ( GraphApprovalStatus.APPROVED if decision["action"] == "approve" else GraphApprovalStatus.REJECTED ) approval.save(update_fields=["status", "updated_at"]) def commit_chapter(self, state): return {"canon_snapshot_id": "canon-1"} def publish_story(self, state): return "test.epub" def test_story_graph_runs_one_consolidated_review() -> None: from langgraph.checkpoint.memory import MemorySaver revision = story_revision() services = FakeStoryServices() runner = StoryWorkflowRunner(build_story_workflow(services, MemorySaver())) graph_run = runner.start(revision) assert graph_run.status == GraphRunStatus.PAUSED assert graph_run.current_node == "approve_plan" assert GraphApproval.objects.filter(status=GraphApprovalStatus.PENDING).count() == 1 graph_run = runner.resume(graph_run.id, {"action": "approve", "actor": "test"}) assert graph_run.status == GraphRunStatus.PAUSED assert graph_run.current_node == "approve_chapter" assert services.audit_calls == 1 assert services.review_calls == [] graph_run = runner.resume(graph_run.id, {"action": "approve", "actor": "test"}) assert graph_run.status == GraphRunStatus.COMPLETE assert graph_run.metadata["final_state"]["export_uri"] == "test.epub" def test_story_graph_allows_one_patch_then_one_final_extraction() -> None: from langgraph.checkpoint.memory import MemorySaver class PatchServices(FakeStoryServices): def __init__(self) -> None: super().__init__() self.extraction_calls = 0 self.patch_calls = 0 self.verification_calls = 0 def extract_continuity(self, state): self.extraction_calls += 1 return {} def judge_state_contract(self, state): self.audit_calls += 1 return {"state_judge_status": "revise", "editorial_finding_ids": ["finding-1"]} def decide_patch(self, state): return {"patch_decision": "patch", "patch_finding_ids": ["finding-1"]} def apply_automatic_patch(self, state): self.patch_calls += 1 return { "revision_id": state["revision_id"], "patch_attempted": True, "patch_status": "applied", "patch_source_revision_id": state["revision_id"], "changed_passages": [{"old_text": "old", "new_text": "new"}], } def verify_patch(self, state): self.verification_calls += 1 return {"verification_status": "pass"} revision = story_revision() services = PatchServices() runner = StoryWorkflowRunner(build_story_workflow(services, MemorySaver())) graph_run = runner.start(revision) graph_run = runner.resume(graph_run.id, {"action": "approve", "actor": "test"}) assert graph_run.current_node == "approve_chapter" assert services.audit_calls == 1 assert services.extraction_calls == 1 assert services.patch_calls == 1 assert services.verification_calls == 1 graph_run = runner.resume( graph_run.id, {"action": "request_revision", "actor": "test"} ) assert graph_run.status == GraphRunStatus.COMPLETE assert services.patch_calls == 1 assert services.verification_calls == 1 def test_patch_selection_excludes_subjective_low_and_optional_contract_findings() -> None: from control_plane.authoring.services import DjangoStoryWorkflowServices revision = story_revision() EditorialFinding.objects.create( revision=revision, review_kind="story_audit", severity="LOW", category="style", description="Optional preference", evidence={"objective": True, "exact_patch_suitable": True}, ) EditorialFinding.objects.create( revision=revision, review_kind="state_contract", severity="MEDIUM", category="contract:BEAT", description="Optional beat detail", evidence={ "objective": True, "exact_patch_suitable": True, "blocking": False, }, ) objective = EditorialFinding.objects.create( revision=revision, review_kind="story_audit", severity="MEDIUM", category="continuity", description="Concrete contradiction", evidence={"objective": True, "exact_patch_suitable": True}, ) services = DjangoStoryWorkflowServices(ModelRouter({})) result = services.decide_patch({"revision_id": str(revision.id)}) assert result["patch_decision"] == "patch" assert result["patch_finding_ids"] == [str(objective.id)] def test_approval_inbox_resumes_story_checkpoint(monkeypatch) -> None: from langgraph.checkpoint.memory import MemorySaver from control_plane.projects import views revision = story_revision() services = FakeStoryServices() checkpointer = MemorySaver() graph_run = StoryWorkflowRunner(build_story_workflow(services, checkpointer)).start(revision) approval = GraphApproval.objects.get(graph_run=graph_run, status=GraphApprovalStatus.PENDING) assert graph_run.execution_graph_version.graph.name == "story_authoring" @contextmanager def checkpointer_context(): yield checkpointer monkeypatch.setattr(views, "open_story_checkpointer", checkpointer_context) monkeypatch.setattr(views, "DjangoStoryWorkflowServices", lambda *_args, **_kwargs: services) response = Client().post( reverse("approval_action", args=[approval.id]), {"action": "approve", "notes": "Keep the opening quiet."}, ) graph_run.refresh_from_db() approval.refresh_from_db() assert response.status_code == 302 assert approval.status == GraphApprovalStatus.APPROVED, list( GraphApproval.objects.filter(graph_run=graph_run).values_list("reason", "status") ) assert graph_run.status == GraphRunStatus.PAUSED assert graph_run.current_node == "approve_chapter" class InterruptedStreamingProvider: provider_name = "test" def __init__(self) -> None: self.calls = 0 def stream(self, request): self.calls += 1 if self.calls == 1: yield ModelChunk("one two three ") raise RuntimeError("connection lost") yield ModelChunk("four five [[END_OF_CHAPTER]]") def health(self): return "AVAILABLE" def test_streamed_story_draft_resumes_preserved_fragment(tmp_path: Path) -> None: provider = InterruptedStreamingProvider() writer = ResumableDraftWriter(ModelRouter({"terra": provider})) result = writer.generate( request=ModelRequestContract( purpose="STORY_PROSE", prompt="write", model_hint="terra" ), partial_path=tmp_path / "chapter.partial.md", minimum_words=5, max_attempts=2, ) assert result.text == "one two three four five" assert result.resumed is False assert provider.calls == 2 assert not (tmp_path / "chapter.partial.md.attempt").exists() def test_short_completed_draft_restarts_without_duplication(tmp_path: Path) -> None: class ShortThenCompleteProvider: provider_name = "test" def __init__(self) -> None: self.calls = 0 def stream(self, request): self.calls += 1 if self.calls == 1: yield ModelChunk("too short [[END_OF_CHAPTER]]") else: yield ModelChunk("one two three four five [[END_OF_CHAPTER]]") def health(self): return "AVAILABLE" provider = ShortThenCompleteProvider() writer = ResumableDraftWriter(ModelRouter({"qwen": provider})) result = writer.generate( request=ModelRequestContract( purpose="STORY_PROSE", prompt="write", model_hint="qwen" ), partial_path=tmp_path / "chapter.partial.md", minimum_words=5, max_attempts=2, ) assert result.text == "one two three four five" assert provider.calls == 2 def test_short_completed_draft_stops_after_one_regeneration(tmp_path: Path) -> None: class AlwaysShortProvider: provider_name = "test" def __init__(self) -> None: self.calls = 0 def stream(self, request): self.calls += 1 yield ModelChunk("too short [[END_OF_CHAPTER]]") def health(self): return "AVAILABLE" provider = AlwaysShortProvider() writer = ResumableDraftWriter(ModelRouter({"qwen": provider})) with pytest.raises(RuntimeError, match="shorter than 5 words"): writer.generate( request=ModelRequestContract( purpose="STORY_PROSE", prompt="write", model_hint="qwen" ), partial_path=tmp_path / "chapter.partial.md", minimum_words=5, max_attempts=4, ) assert provider.calls == 2 def test_draft_chapter_generates_full_chapter_in_one_call(tmp_path: Path) -> None: revision = story_revision() revision.chapter.story.artifact_root = str(tmp_path) revision.chapter.story.save(update_fields=["artifact_root", "updated_at"]) revision.scene_plan = { "target_words": 2000, "scenes": [ {"number": 1, "word_budget": 1000, "beats": ["first"]}, {"number": 2, "word_budget": 1000, "beats": ["second"]}, ], } revision.save(update_fields=["scene_plan", "updated_at"]) services = DjangoStoryWorkflowServices(ModelRouter({})) services._context = lambda state, current: { "chapter": {}, "structured_canon": {}, "prior_canon": {}, "previous_chapter_tail": "", } class ChapterWriter: def __init__(self) -> None: self.requests = [] def generate(self, *, request, **kwargs): self.requests.append(request) assert kwargs["max_attempts"] == 1 return DraftResult("Complete chapter prose.", 1, False, 3) writer = ChapterWriter() services.writer = writer services.draft_chapter({"revision_id": str(revision.id)}) revision.refresh_from_db() assert revision.prose == "Complete chapter prose." assert len(writer.requests) == 1 assert writer.requests[0].model_hint == "terra" assert revision.generation_metadata["draft_mode"] == "full_chapter" def test_compact_scene_contract_consolidates_all_required_beats_into_three() -> None: scene = { "number": 1, "purpose": "Test", "beats": [{"text": f"beat {number}", "required": True} for number in range(1, 13)], } compact = compact_scene_contract(scene) assert len(compact["beats"]) == 3 assert [beat["source_beat_count"] for beat in compact["beats"]] == [4, 4, 4] combined = " ".join(beat["text"] for beat in compact["beats"]) assert all(f"beat {number}" in combined for number in range(1, 13)) def test_scene_draft_packet_excludes_other_scenes() -> None: selected = {"number": 1, "beats": [{"text": "selected", "required": True}]} packet = scene_draft_packet( { "target_words": 2000, "exact_values": ["exact"], "scenes": [selected, {"number": 2, "beats": [{"text": "unrelated"}]}], }, selected, ) assert packet["target_words"] == 1000 assert packet["scene"]["beats"] == [{"text": "selected", "required": True}] assert "scenes" not in packet["chapter_scope"] def test_compact_chapter_plan_preserves_all_required_beats() -> None: plan = { "target_words": 5000, "scenes": [ { "number": number, "beats": [ {"text": f"scene {number} beat {beat}", "required": True} for beat in range(1, 7) ], } for number in range(1, 3) ], } compact = compact_chapter_plan(plan) assert [len(scene["beats"]) for scene in compact["scenes"]] == [3, 3] combined = json.dumps(compact) assert all( f"scene {scene} beat {beat}" in combined for scene in range(1, 3) for beat in range(1, 7) ) def test_deterministic_temporal_findings_catch_premature_wealth() -> None: prose = ( "The dungeon had not noticed that he had become wealthy.\n\n" "The bids were opened one at a time." ) findings = deterministic_temporal_findings(prose) assert len(findings) == 1 assert findings[0]["evidence_quote"] == "The dungeon had not noticed that he had become wealthy." assert findings[0]["suggested_revision"] == "The dungeon had not noticed that he might soon become wealthy." def test_deterministic_temporal_findings_catch_exact_payout_before_bidding() -> None: prose = ( "Three and a half million crowns would buy better weapons.\n\n" "The bids were opened one at a time." ) plan = { "exact_values": [ "Corin's five-sixths finder share is exactly 3,500,000 silver crowns." ] } findings = deterministic_temporal_findings(prose, plan) assert len(findings) == 1 assert findings[0]["category"] == "premature_exact_value" assert findings[0]["suggested_revision"] == "The sale proceeds would buy better weapons." def test_epub_contains_approved_chapter_entries(tmp_path: Path) -> None: destination = write_epub( title="Test Book", series="Test Series", chapters=[{"title": "Chapter 1: Opening", "content": "First paragraph.\n\nSecond."}], destination=tmp_path / "book.epub", ) with zipfile.ZipFile(destination) as archive: content = archive.read("OEBPS/chapter-1.xhtml").decode("utf-8") assert "Chapter 1: Opening" in content assert "First paragraph." in content def test_story_versions_hash_content() -> None: revision = story_revision() assert len(revision.story_bible.sha256) == 64 assert len(revision.outline.sha256) == 64 def test_exact_patch_preserves_unaffected_prose() -> None: prose = "First paragraph.\n\nThe water remained.\n\nLast paragraph." revised = apply_exact_edits( prose, [{"old_text": "The water remained.", "new_text": "The water ran thinner than before."}], ) assert revised == "First paragraph.\n\nThe water ran thinner than before.\n\nLast paragraph." def test_exact_patch_rejects_ambiguous_source_text() -> None: import pytest with pytest.raises(ValueError, match="exactly once"): apply_exact_edits( "Repeated. Repeated.", [{"old_text": "Repeated.", "new_text": "Changed."}], ) def test_contract_requirements_receive_stable_ids() -> None: requirements = build_contract_requirements( { "day_start": "Morning", "day_end": "Evening", "scenes": [ { "number": 2, "beats": ["First beat", "Second beat"], "ending_state": "The door is closed.", } ], "exact_values": ["Ten crowns"], "forbidden_shortcuts": ["Do not montage the sale."], } ) assert [item["id"] for item in requirements] == [ "S02-B01", "S02-B02", "S02-END", "VALUE-01", "SHORTCUT-01", "TIME-START", "TIME-END", ] assert requirements[0]["blocking"] is False assert requirements[3]["blocking"] is True assert requirements[4]["blocking"] is False def test_contract_marks_only_explicit_required_beats_as_blocking() -> None: requirements = build_contract_requirements( { "scenes": [ { "number": 1, "beats": [ {"text": "The sale settles.", "required": True}, {"text": "Rain taps the window.", "required": False}, ], } ] } ) assert requirements[0]["blocking"] is True assert requirements[1]["blocking"] is False def test_contract_caps_required_beats_per_scene() -> None: requirements = build_contract_requirements( { "scenes": [ { "number": 1, "beats": [ {"text": f"Beat {index}", "required": True} for index in range(7) ], } ] } ) assert sum(1 for item in requirements if item["blocking"]) == 3 def test_exact_patch_rejects_overlapping_and_over_budget_edits() -> None: import pytest with pytest.raises(ValueError, match="overlap"): apply_exact_edits( "abcdefghij", [ {"old_text": "abcde", "new_text": "ABCDE"}, {"old_text": "defgh", "new_text": "DEFGH"}, ], ) with pytest.raises(ValueError, match="limit"): apply_exact_edits( "a" * 100 + " target " + "b" * 100, [{"old_text": " target ", "new_text": " a much longer replacement passage "}], max_change_ratio=0.05, ) def test_state_changes_build_queryable_book_snapshot() -> None: state = apply_state_changes( {}, [ { "sequence": 1, "entity_key": "character.corin.vale", "entity_kind": "character", "canonical_name": "Corin Vale", "predicate": "finances.balance", "operation": "SET", "previous_value": None, "new_value": 10, } ], through_chapter=2, chapter_state={"summary": "Corin receives ten crowns."}, ) assert state["through_chapter"] == 2 assert state["entities"]["character.corin.vale"]["facts"]["finances"]["balance"] == 10 def test_transfer_state_change_records_explicit_custody_destination() -> None: state = apply_state_changes( {}, [ { "sequence": 1, "entity_key": "object.envelope.after", "entity_kind": "object", "canonical_name": "Envelope After", "predicate": "custody.holder", "operation": "TRANSFER", "previous_value": None, "new_value": None, "related_entity_key": "character.sabine", } ], through_chapter=2, chapter_state={}, ) envelope = state["entities"]["object.envelope.after"] assert envelope["facts"]["custody"]["holder"] == "character.sabine" assert envelope["relations"]["custody.holder"] == "character.sabine" def test_transfer_state_change_requires_destination_entity() -> None: with pytest.raises(ValueError, match="related_entity_key"): apply_state_changes( {}, [ { "sequence": 1, "entity_key": "object.envelope.after", "predicate": "custody.holder", "operation": "TRANSFER", "previous_value": None, "new_value": None, } ], through_chapter=2, chapter_state={}, ) class StateWorkflowProvider: provider_name = "test" def complete(self, request): if str(request.purpose) == "STORY_CONTINUITY": content = """{ "schema_version": 2, "through_chapter": 1, "state_document": { "timeline": {"relative": "Morning"}, "scene_end": {"location": "Market"}, "characters": {"Corin Vale": {"location": "Market"}}, "inventory": [], "money": [{"owner": "Corin Vale", "balance": 10}], "relationships": [], "open_threads": [], "promises_and_constraints": [], "reveals": {}, "chapter_summary": ["Corin received ten crowns."] }, "changes": [{ "entity_key": "character.corin.vale", "entity_kind": "character", "canonical_name": "Corin Vale", "change_type": "MONEY_CHANGED", "predicate": "finances.balance", "operation": "SET", "previous_value": null, "new_value": 10, "related_entity_key": "", "evidence_quote": "Corin received ten crowns.", "evidence_location": "Scene 1" }] }""" else: content = """{ "verdict": "PASS", "requirements": [ {"requirement_id":"S01-B01","status":"HIT","evidence_quote":"Corin received ten crowns.","evidence_location":"Scene 1","details":""}, {"requirement_id":"S01-END","status":"HIT","evidence_quote":"Corin received ten crowns.","evidence_location":"Scene 1","details":""} ] }""" return ModelResponseContract(model="test", content=content, metadata={}) def test_validated_state_document_commits_entity_history(tmp_path: Path) -> None: from control_plane.authoring.services import DjangoStoryWorkflowServices revision = story_revision() story = revision.chapter.story story.artifact_root = str(tmp_path) story.save(update_fields=["artifact_root", "updated_at"]) revision.scene_plan = { "scenes": [ { "number": 1, "beats": ["Corin receives ten crowns."], "ending_state": "Corin has ten crowns.", } ] } revision.prose = "Corin received ten crowns." revision.save() context = { "chapter": {"number": 1, "title": "Opening", "outline": {}}, "story_bible": "Canon", "structured_canon": {}, "prior_canon": {}, "previous_chapter_tail": "", "source_revision": "", } snapshot = GenerationContextSnapshot.objects.create( story=story, chapter=revision.chapter, story_bible=revision.story_bible, outline=revision.outline, content=context, sha256="0" * 64, ) revision.context_snapshot = snapshot revision.save(update_fields=["context_snapshot", "updated_at"]) provider = StateWorkflowProvider() services = DjangoStoryWorkflowServices( ModelRouter({"qwen": provider, "terra": provider, "luna": provider}) ) state = { "revision_id": str(revision.id), "story_id": str(story.id), "context_snapshot_id": str(snapshot.id), } services.extract_continuity(state) result = services.judge_state_contract(state) commit = services.commit_chapter(state) document = ChapterStateDocument.objects.get(revision=revision) change = StateChange.objects.get(state_document=document) canon = CanonSnapshot.objects.get(id=commit["canon_snapshot_id"]) assert result["state_judge_status"] == "pass" assert document.status == "COMMITTED" assert StoryEntity.objects.get(entity_key="character.corin.vale").canonical_name == "Corin Vale" assert change.status == StateChangeStatus.COMMITTED assert canon.state["entities"]["character.corin.vale"]["facts"]["finances"]["balance"] == 10 assert Path(document.json_artifact_uri).exists() assert Path(document.markdown_artifact_uri).exists() def test_evidence_accepts_ordered_dialogue_fragments_with_attribution() -> None: prose = '"Replace those clothes," she said. "Eat somewhere respectable. Make no promises tonight."' assert evidence_is_present( prose, "Replace those clothes. Eat somewhere respectable. Make no promises tonight.", )