Merge pull request #105 from 7df-lab/dev/goal_title_fix_0616

Fix goal session title generation

Author: wangtsiao

crates/server/src/runtime/goal_handlers.rs

@@ -20,6 +20,14 @@ impl ServerRuntime {
         };
         let session_id = params.session_id;
         let replace_existing = params.replace_existing;
+        let title_input = params.objective.trim().to_string();
+        if !self.sessions.lock().await.contains_key(&session_id) {
+            return self.error_response(
+                request_id,
+                ProtocolErrorCode::SessionNotFound,
+                "session does not exist",
+            );
+        }
 
         let mut stores = self.goal_stores.lock().await;
         let store = stores.entry(session_id).or_insert_with(GoalStore::new);
@@ -43,6 +51,8 @@ impl ServerRuntime {
                     tracing::warn!(session_id = %session_id, error = %error, "failed to persist goal create record");
                 }
                 self.sync_core_session_goal(session_id, session_goal).await;
+                self.maybe_start_title_generation_from_user_input(session_id, &title_input)
+                    .await;
                 if replace_existing {
                     self.interrupt_active_goal_continuation_turn(session_id, "goal replaced")
                         .await;
@@ -77,10 +87,23 @@ impl ServerRuntime {
         };
         let session_id = params.session_id;
         let requested_status = params.status;
+        let title_input = params
+            .objective
+            .as_deref()
+            .map(str::trim)
+            .filter(|objective| !objective.is_empty())
+            .map(str::to_string);
         let only_pause_budget_limited = requested_status
             == Some(devo_protocol::ThreadGoalStatus::Paused)
             && params.objective.is_none()
             && params.token_budget.is_none();
+        if !self.sessions.lock().await.contains_key(&session_id) {
+            return self.error_response(
+                request_id,
+                ProtocolErrorCode::SessionNotFound,
+                "session does not exist",
+            );
+        }
 
         let mut stores = self.goal_stores.lock().await;
         let store = stores.entry(session_id).or_insert_with(GoalStore::new);
@@ -146,6 +169,10 @@ impl ServerRuntime {
                     )
                     .await;
                 }
+                if let Some(title_input) = title_input {
+                    self.maybe_start_title_generation_from_user_input(session_id, &title_input)
+                        .await;
+                }
                 if should_continue {
                     self.maybe_start_goal_continuation_turn(session_id).await;
                 }

crates/server/src/runtime/handlers/turn.rs

@@ -304,32 +304,8 @@ impl ServerRuntime {
             )
             .await;
         }
-        self.maybe_assign_provisional_title(params.session_id, &display_input)
+        self.maybe_start_title_generation_from_user_input(params.session_id, &display_input)
             .await;
-        {
-            let mut session = session_arc.lock().await;
-            if session.first_user_input.is_none() {
-                session.first_user_input = Some(display_input.clone());
-            }
-        }
-        let needs_title = {
-            let session = session_arc.lock().await;
-            let first_input = session.first_user_input.clone();
-            let needs = matches!(
-                session.summary.title_state,
-                SessionTitleState::Unset | SessionTitleState::Provisional
-            );
-            (needs, first_input)
-        };
-        if needs_title.0
-            && let Some(first_input) = needs_title.1
-        {
-            let runtime = Arc::clone(self);
-            let sid = params.session_id;
-            tokio::spawn(async move {
-                runtime.maybe_generate_final_title(sid, first_input).await;
-            });
-        }
         let (record, session_context, turn_context) = {
             let session = session_arc.lock().await;
             let core_session = session.core_session.lock().await;

crates/server/src/runtime/items.rs

@@ -1,6 +1,41 @@
 use super::*;
 
 impl ServerRuntime {
+    pub(super) async fn maybe_start_title_generation_from_user_input(
+        self: &Arc<Self>,
+        session_id: SessionId,
+        user_input: &str,
+    ) {
+        self.maybe_assign_provisional_title(session_id, user_input)
+            .await;
+
+        let needs_title = {
+            let Some(session_arc) = self.sessions.lock().await.get(&session_id).cloned() else {
+                return;
+            };
+            let mut session = session_arc.lock().await;
+            if session.first_user_input.is_none() {
+                session.first_user_input = Some(user_input.to_string());
+            }
+            let first_input = session.first_user_input.clone();
+            let needs = matches!(
+                session.summary.title_state,
+                SessionTitleState::Unset | SessionTitleState::Provisional
+            );
+            (needs, first_input)
+        };
+        if needs_title.0
+            && let Some(first_input) = needs_title.1
+        {
+            let runtime = Arc::clone(self);
+            tokio::spawn(async move {
+                runtime
+                    .maybe_generate_final_title(session_id, first_input)
+                    .await;
+            });
+        }
+    }
+
     pub(super) async fn maybe_assign_provisional_title(
         &self,
         session_id: SessionId,