fix puzzle generation failure stage

server-rs/crates/spacetime-module/src/puzzle.rs

@@ -1042,6 +1042,7 @@ fn mark_puzzle_draft_generation_failed_tx(
             ))
         }
     };
+    let next_stage = resolve_failed_puzzle_agent_stage(row.stage, &draft);
     upsert_puzzle_draft_work_profile(
         ctx,
         &row.session_id,
@@ -1059,7 +1060,7 @@ fn mark_puzzle_draft_generation_failed_tx(
             seed_text: row.seed_text.clone(),
             current_turn: row.current_turn,
             progress_percent: row.progress_percent.max(88),
-            stage: row.stage,
+            stage: next_stage,
             anchor_pack_json: row.anchor_pack_json.clone(),
             draft_json: Some(serialize_json(&draft)),
             last_assistant_reply: Some(input.error_message),
@@ -1078,6 +1079,21 @@ fn mark_puzzle_draft_generation_failed_tx(
     )
 }
 
+fn resolve_failed_puzzle_agent_stage(
+    current_stage: PuzzleAgentStage,
+    draft: &PuzzleResultDraft,
+) -> PuzzleAgentStage {
+    if current_stage == PuzzleAgentStage::Published {
+        return PuzzleAgentStage::Published;
+    }
+
+    if build_result_preview(draft, Some("陶泥儿主")).publish_ready {
+        PuzzleAgentStage::ReadyToPublish
+    } else {
+        PuzzleAgentStage::DraftReady
+    }
+}
+
 fn save_puzzle_form_draft_tx(
     ctx: &TxContext,
     input: PuzzleFormDraftSaveInput,
@@ -3862,6 +3878,25 @@ mod tests {
         assert_eq!(draft.work_title, "我的猫街合集");
     }
 
+    #[test]
+    fn failed_generation_returns_result_draft_stage_instead_of_refining() {
+        let anchor_pack = infer_anchor_pack("画面描述：一只猫在雨夜灯牌下回头。", None);
+        let draft = compile_result_draft_from_seed(
+            &anchor_pack,
+            &[],
+            Some("画面描述：一只猫在雨夜灯牌下回头。"),
+        );
+
+        assert_eq!(
+            resolve_failed_puzzle_agent_stage(PuzzleAgentStage::ImageRefining, &draft),
+            PuzzleAgentStage::DraftReady
+        );
+        assert_eq!(
+            resolve_failed_puzzle_agent_stage(PuzzleAgentStage::Published, &draft),
+            PuzzleAgentStage::Published
+        );
+    }
+
     #[test]
     fn puzzle_recommendation_score_prefers_same_author_weight() {
         let left = PuzzleWorkProfile {