Add generation harness runtime

2026-06-21 22:31:38 +08:00
parent 7ebdfb2582
commit 459ca9edef
18 changed files with 2846 additions and 419 deletions
--- a/backend/tests/test_harness_runtime.py
+++ b/backend/tests/test_harness_runtime.py
@@ -0,0 +1,374 @@
+"""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",
+    ]