"""Tests for generation harness runtime support.""" import pytest from sqlalchemy import select from app.db.models import GenerationJob, GenerationJobEvent from app.services.adapters.storybook.primary import Storybook, StorybookPage from app.services.adapters.text.models import StoryOutput from app.services.generation_jobs import create_generation_job, record_generation_event from app.services.harness.control import ExecutionControl, GenerationJobCanceledError from app.services.harness.plans import ( build_asset_plan, build_story_plan, build_storybook_plan, ) from app.services.harness.quality_gates import ( QualityGateError, validate_story_output, validate_storybook_output, ) from app.services.harness.trace import TraceRecorder from app.services.harness.types import ( ArtifactKind, FailureCategory, WorkflowStep, artifact_for_event, normalize_trace_metadata, step_for_event, ) def test_event_type_maps_to_standard_workflow_step(): assert step_for_event("request_accepted") == WorkflowStep.REQUEST_ACCEPTANCE assert step_for_event("context_prepared") == WorkflowStep.CONTEXT_PREPARATION assert step_for_event("narrative_generated") == WorkflowStep.NARRATIVE_GENERATION assert step_for_event("story_saved") == WorkflowStep.STORY_PERSISTENCE assert step_for_event("provider_call_succeeded") == WorkflowStep.PROVIDER_INVOCATION assert step_for_event("quality_gate_failed") == WorkflowStep.NARRATIVE_GENERATION assert step_for_event("cover_image_failed") == WorkflowStep.IMAGE_GENERATION assert step_for_event("audio_succeeded") == WorkflowStep.AUDIO_GENERATION assert step_for_event("generation_canceled") == WorkflowStep.CANCELLATION assert step_for_event("generation_stale_failed") == WorkflowStep.STALE_RECOVERY assert step_for_event("future_event") == WorkflowStep.UNKNOWN def test_event_type_maps_to_standard_artifact(): assert artifact_for_event("narrative_generated") == ArtifactKind.STORY_TEXT assert artifact_for_event("quality_gate_failed") == ArtifactKind.STORY_TEXT assert artifact_for_event("cover_image_succeeded") == ArtifactKind.COVER_IMAGE assert artifact_for_event("storybook_page_image_failed") == ArtifactKind.PAGE_IMAGE assert artifact_for_event("audio_cache_hit") == ArtifactKind.AUDIO assert artifact_for_event("postprocessing_queued") == ArtifactKind.ACHIEVEMENT_MEMORY assert artifact_for_event("request_accepted") == ArtifactKind.NONE def test_trace_metadata_adds_standard_fields_without_dropping_legacy_values(): metadata = normalize_trace_metadata( "provider_call_failed", { "capability": "text", "adapter": "demo", "error": "timeout", }, failure_category=FailureCategory.TIMEOUT, retryable=True, ) assert metadata["capability"] == "text" assert metadata["adapter"] == "demo" assert metadata["error"] == "timeout" assert metadata["step"] == "provider_invocation" assert metadata["artifact"] == "none" assert metadata["failure_category"] == "timeout" assert metadata["retryable"] is True assert metadata["blocks_main_result"] is False def test_trace_metadata_respects_explicit_step_and_artifact(): metadata = normalize_trace_metadata( "narrative_generated", {"title": "小兔子的冒险"}, step=WorkflowStep.NARRATIVE_GENERATION, artifact=ArtifactKind.STORYBOOK_PAGES, blocks_main_result=True, ) assert metadata["title"] == "小兔子的冒险" assert metadata["step"] == "narrative_generation" assert metadata["artifact"] == "storybook_pages" assert metadata["blocks_main_result"] is True def test_story_plan_without_assets_snapshot(): assert build_story_plan(generate_images=False).to_snapshot() == { "mode": "story", "tasks": [ { "key": "prepare_context", "step": "context_preparation", "artifact": "none", "required": True, "recoverable": False, }, { "key": "generate_narrative", "step": "narrative_generation", "artifact": "story_text", "required": True, "recoverable": False, }, { "key": "persist_story", "step": "story_persistence", "artifact": "story_text", "required": True, "recoverable": False, }, { "key": "queue_postprocessing", "step": "postprocessing", "artifact": "achievement_memory", "required": False, "recoverable": True, }, { "key": "complete_generation", "step": "completion", "artifact": "none", "required": True, "recoverable": False, }, ], } def test_story_plan_with_assets_marks_cover_recoverable(): plan = build_story_plan(generate_images=True).to_snapshot() assert plan["mode"] == "story_with_assets" assert plan["tasks"][3] == { "key": "generate_cover_image", "step": "image_generation", "artifact": "cover_image", "required": False, "recoverable": True, } def test_storybook_plan_with_images_marks_storybook_images_recoverable(): plan = build_storybook_plan(generate_images=True).to_snapshot() assert plan["mode"] == "storybook" assert [task["key"] for task in plan["tasks"]] == [ "prepare_context", "generate_storybook_pages", "generate_storybook_images", "persist_storybook", "queue_postprocessing", "complete_generation", ] assert plan["tasks"][2]["artifact"] == "image" assert plan["tasks"][2]["recoverable"] is True def test_asset_retry_plan_deduplicates_assets(): plan = build_asset_plan(output_mode="asset_retry", assets=["image", "audio", "image"]) assert plan.to_snapshot() == { "mode": "asset_retry", "tasks": [ { "key": "start_asset_retry", "step": "asset_retry", "artifact": "none", "required": True, "recoverable": False, }, { "key": "complete_image_asset", "step": "image_generation", "artifact": "image", "required": False, "recoverable": True, }, { "key": "complete_audio_asset", "step": "audio_generation", "artifact": "audio", "required": False, "recoverable": True, }, { "key": "complete_asset_retry", "step": "asset_retry", "artifact": "none", "required": True, "recoverable": False, }, ], } def test_story_quality_gate_accepts_complete_child_safe_story(): validate_story_output( StoryOutput( mode="generated", title="小兔子的月光花园", story_text="小兔子在花园里学会了和朋友轮流分享水壶。", cover_prompt_suggestion="A gentle moonlit garden with a rabbit", ) ) def test_story_quality_gate_rejects_missing_story_text(): output = StoryOutput( mode="generated", title="空白故事", story_text="", cover_prompt_suggestion="A cover", ) try: validate_story_output(output) except QualityGateError as exc: assert [issue.code.value for issue in exc.issues] == ["missing_story_text"] assert exc.to_metadata()["issues"][0]["field"] == "story_text" else: raise AssertionError("Expected QualityGateError") def test_story_quality_gate_rejects_obviously_unsafe_child_content(): output = StoryOutput( mode="generated", title="危险词测试", story_text="这个故事包含血腥场景。", cover_prompt_suggestion="A cover", ) try: validate_story_output(output) except QualityGateError as exc: assert [issue.code.value for issue in exc.issues] == ["unsafe_child_content"] assert exc.to_metadata()["issues"][0]["failure_category"] == "safety_error" else: raise AssertionError("Expected QualityGateError") def test_storybook_quality_gate_rejects_duplicate_page_number(): storybook = Storybook( title="森林绘本", main_character="小兔子", art_style="水彩", cover_prompt="A forest cover", pages=[ StorybookPage(page_number=1, text="第一页。", image_prompt="page 1"), StorybookPage(page_number=1, text="第二页。", image_prompt="page 2"), ], ) try: validate_storybook_output(storybook) except QualityGateError as exc: assert [issue.code.value for issue in exc.issues] == [ "invalid_storybook_page_number" ] assert exc.to_metadata()["issues"][0]["field"] == "pages[1].page_number" else: raise AssertionError("Expected QualityGateError") @pytest.mark.asyncio async def test_trace_recorder_persists_standard_metadata(db_session, test_user): job = await create_generation_job( db_session, user_id=test_user.id, output_mode="story", input_type="keywords", request_payload={"data": "小兔子"}, ) event = await TraceRecorder(db_session).record_step( job=job, event_type="provider_call_failed", status="failed", metadata={ "capability": "text", "adapter": "demo", "error": "timeout", }, failure_category=FailureCategory.TIMEOUT, retryable=True, ) assert event is not None events = ( await db_session.execute( select(GenerationJobEvent) .where(GenerationJobEvent.job_id == job.id) .order_by(GenerationJobEvent.id) ) ).scalars().all() assert [item.event_type for item in events] == [ "request_accepted", "provider_call_failed", ] metadata = events[1].event_metadata assert metadata["capability"] == "text" assert metadata["adapter"] == "demo" assert metadata["step"] == "provider_invocation" assert metadata["artifact"] == "none" assert metadata["failure_category"] == "timeout" assert metadata["retryable"] is True @pytest.mark.asyncio async def test_trace_recorder_ignores_missing_job(db_session): event = await TraceRecorder(db_session).record_step( job=None, event_type="context_prepared", status="succeeded", ) assert event is None @pytest.mark.asyncio async def test_execution_control_cancels_job_at_safe_checkpoint( db_session, test_user, test_story, ): job = await create_generation_job( db_session, user_id=test_user.id, output_mode="story", input_type="keywords", request_payload={"data": "小兔子"}, story_id=test_story.id, ) await record_generation_event( db_session, job=job, story_id=test_story.id, event_type="cancel_requested", status="running", message="Cancellation requested.", ) with pytest.raises(GenerationJobCanceledError): await ExecutionControl(db_session).stop_if_cancel_requested( job=job, story=test_story, ) refreshed_job = ( await db_session.execute(select(GenerationJob).where(GenerationJob.id == job.id)) ).scalar_one() assert refreshed_job.status == "canceled" assert refreshed_job.current_step == "generation_canceled" assert refreshed_job.error_message == "Generation canceled by user." events = ( await db_session.execute( select(GenerationJobEvent) .where(GenerationJobEvent.job_id == job.id) .order_by(GenerationJobEvent.id) ) ).scalars().all() assert [item.event_type for item in events] == [ "request_accepted", "cancel_requested", "generation_canceled", ]