1

server-rs/crates/api-server/src/puzzle.rs

@@ -52,6 +52,7 @@ use spacetime_client::{
 use std::convert::Infallible;
 
 use crate::{
+    ai_generation_drafts::{AiGenerationDraftContext, AiGenerationDraftWriter},
     api_response::json_success_body,
     auth::AuthenticatedAccessToken,
     http_error::AppError,
@@ -288,6 +289,16 @@ pub async fn stream_puzzle_agent_message(
     let session_id_for_stream = session_id.clone();
     let owner_user_id_for_stream = owner_user_id.clone();
     let stream = async_stream::stream! {
+        let mut draft_writer = AiGenerationDraftWriter::new(AiGenerationDraftContext::new(
+            "puzzle",
+            owner_user_id_for_stream.as_str(),
+            session_id_for_stream.as_str(),
+            payload.client_message_id.as_str(),
+            "拼图模板生成草稿",
+        ));
+        if let Err(error) = draft_writer.ensure_started(state.spacetime_client()).await {
+            tracing::warn!(error = %error, "拼图模板生成草稿任务启动失败，主生成流程继续执行");
+        }
         let (reply_tx, mut reply_rx) = tokio::sync::mpsc::unbounded_channel::<String>();
         let turn_result = {
             let run_turn = run_puzzle_agent_turn(
@@ -306,6 +317,7 @@ pub async fn stream_puzzle_agent_message(
                     result = &mut run_turn => break result,
                     maybe_text = reply_rx.recv() => {
                         if let Some(text) = maybe_text {
+                            draft_writer.persist_visible_text(state.spacetime_client(), text.as_str()).await;
                             yield Ok::<Event, Infallible>(puzzle_sse_json_event_or_error(
                                 "reply_delta",
                                 json!({ "text": text }),
@@ -317,6 +329,7 @@ pub async fn stream_puzzle_agent_message(
         };
 
         while let Some(text) = reply_rx.recv().await {
+            draft_writer.persist_visible_text(state.spacetime_client(), text.as_str()).await;
             yield Ok::<Event, Infallible>(puzzle_sse_json_event_or_error(
                 "reply_delta",
                 json!({ "text": text }),