1

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

@@ -42,8 +42,14 @@ use crate::big_fish_agent_turn::{
     run_big_fish_agent_turn,
 };
 use crate::{
-    api_response::json_success_body, auth::AuthenticatedAccessToken, http_error::AppError,
-    request_context::RequestContext, state::AppState,
+    ai_generation_drafts::{
+        AiGenerationDraftContext, AiGenerationDraftSink, AiGenerationDraftWriter,
+    },
+    api_response::json_success_body,
+    auth::AuthenticatedAccessToken,
+    http_error::AppError,
+    request_context::RequestContext,
+    state::AppState,
 };
 
 pub async fn create_big_fish_session(
@@ -203,14 +209,44 @@ pub async fn submit_big_fish_message(
         .map_err(|error| {
             big_fish_error_response(&request_context, map_big_fish_client_error(error))
         })?;
+    let mut draft_writer = AiGenerationDraftWriter::new(AiGenerationDraftContext::new(
+        "big_fish",
+        owner_user_id.as_str(),
+        session_id.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 draft_sink = AiGenerationDraftSink::new(
+        AiGenerationDraftContext::new(
+            "big_fish",
+            owner_user_id.as_str(),
+            session_id.as_str(),
+            payload.client_message_id.as_str(),
+            "大鱼吃小鱼模板生成草稿",
+        ),
+        state.spacetime_client().clone(),
+    );
     let turn_result = run_big_fish_agent_turn(
         BigFishAgentTurnRequest {
             llm_client: state.llm_client(),
             session: &submitted_session,
         },
-        |_| {},
+        move |text| {
+            draft_sink.persist_visible_text_async(text);
+        },
     )
     .await;
+    if let Ok(result) = &turn_result {
+        draft_writer
+            .persist_visible_text(
+                state.spacetime_client(),
+                result.assistant_reply_text.as_str(),
+            )
+            .await;
+    }
     let finalize_input = match turn_result {
         Ok(turn_result) => build_finalize_record_input(
             session_id.clone(),
@@ -285,6 +321,26 @@ pub async fn stream_big_fish_message(
         .map_err(|error| {
             big_fish_error_response(&request_context, map_big_fish_client_error(error))
         })?;
+    let mut draft_writer = AiGenerationDraftWriter::new(AiGenerationDraftContext::new(
+        "big_fish",
+        owner_user_id.as_str(),
+        session_id.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 draft_sink = AiGenerationDraftSink::new(
+        AiGenerationDraftContext::new(
+            "big_fish",
+            owner_user_id.as_str(),
+            session_id.as_str(),
+            payload.client_message_id.as_str(),
+            "大鱼吃小鱼模板生成草稿",
+        ),
+        state.spacetime_client().clone(),
+    );
     let mut streamed_reply_text = String::new();
     let turn_result = run_big_fish_agent_turn(
         BigFishAgentTurnRequest {
@@ -292,10 +348,16 @@ pub async fn stream_big_fish_message(
             session: &submitted_session,
         },
         |text| {
+            draft_sink.persist_visible_text_async(text);
             streamed_reply_text = text.to_string();
         },
     )
     .await;
+    if !streamed_reply_text.is_empty() {
+        draft_writer
+            .persist_visible_text(state.spacetime_client(), streamed_reply_text.as_str())
+            .await;
+    }
     let reply_text = match &turn_result {
         Ok(result) => result.assistant_reply_text.clone(),
         Err(error) => error.to_string(),